Towards Understanding Grokking: An Effective Theory of Representation Learning
Anonymous Author(s)
Abstract
We aim to understand grokking , a phenomenon where models generalize long afteroverfitting their training set.
We present both a microscopic analysis anchored by aneffective theory and a macroscopic analysis of phase diagrams describing learningperformance across hyperparameters.
We find that generalization originates fromstructured representations whose training dynamics and dependence on training setsize can be predicted by our effective theory in a toy setting.
We observe empiricallythe presence of four learning phases: comprehension , grokking , memorization , andconfusion .
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 thememorization phase, leading to delayed generalization.
The Goldilocks phaseis reminiscent of “intelligence from starvation” in Darwinian evolution, whereresource limitations drive discovery of more efficient solutions.
This study not onlyprovides intuitive explanations of the origin of grokking, but also highlights theusefulness of physics-inspired tools, e.g., effective theories and phase diagrams,for understanding deep learning.
Introduction
Perhaps the central challenge of a scientific understanding of deep learning lies in accounting for neural network generalization.
Power et al.
[1] recently added a new puzzle to the task of understandinggeneralization with their discovery of grokking .
Grokking refers to the surprising phenomenon ofdelayed generalization where neural networks, on certain learning problems, generalize long afteroverfitting their training set.
It is a rare albeit striking phenomenon that violates common machinelearning intuitions. We attempt to answer the following three questions:
Q
The origin of generalization : When trained on the algorithmic datasets where grokkingoccurs, how do models generalize at all?
Q
The critical training size : Why does the training time needed to “grok” (generalize) divergeas the training set size decreases toward a critical point?Q3 Delayed generalization : Under what conditions does delayed generalization occur?
We provide evidence that representation learning is central to answering each of these questions.
Ouranswers can be summarized as follows:
A1 Generalization can be attributed to learning a good representation of the inputs,
i e. a representation that has the appropriate structure for the task.
See Figures 1 and 2.
A2 The critical training set size corresponds to the least amount of training data that can deter- mine such a representation (which, in some cases, is unique up to linear transformations). [Figure]
A3 Grokking is a phase between “comprehension” and “memorization” phases and it can beremedied with proper hyperparmeter tuning, as illustrated by the phase diagrams in Figure 5.
This paper is organized as follows:
In Section 2, we introduce the problem setting and build asimplified toy model.
In Section 3, we will use an effective theory approach, a useful tool fromtheoretical physics, to shed some light on questions Q1 and Q2 and show the relationship betweengeneralization and the learning of structured representations.
In Section 4, we explain Q3 bydisplaying 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. {+
+}
Problem Setting
Power et al.
[1] observe grokking on an unusual task – learning an “algorithmic” binary operation.
Given some binary operation ◦ , a network is tasked with learning the map ( a, b ) (cid:55)→ c where a ◦ b = c .
They use a decoder-only transformer to predict the second to last token in a tokenized equation of theform “<lhs> <op> <rhs> <eq> <result> <eos>”.
Each token is represented as a 256-dimensionalembedding vector.
The embeddings are learnable and initialized randomly.
After the transformer, afinal 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 behaviorsfrom the setup of [1].
Although [1] treated this as a classification task, we study both regression andclassification problems.
The basic setup is as follows: our model takes as input the symbols a, b andmaps them to trainable embedding vectors E a , E b ∈ R d in .
It then sums E a , E b and sends the resultingvector through a “decoder” MLP.
The target output vector, denoted Y c ∈ R d out is a fixed randomvector (regression task) or a one-hot vector (classification task).
Our model architecture can thereforebe compactly described as ( a, b ) (cid:55)→
Dec( E a + E b ) , where the embeddings E ∗ and the decoder aretrainable.
Despite its simplicity, this toy model can generalize to all abelian groups (discussed in Appendix A).
In sections 3-4.1, we consider only the binary operation of addition.
We considermodular addition in Section 4.2 to generalize some of our results to a transformer architecture andstudy general non-abelian operations in Appendix F.
Dataset In our toy setting, we are concerned with learning the addition operation.
A data samplecorresponding to i + j is denoted as ( i, j ) for simplicity.
If i, j ∈ { 0 , . . . , p − 1 } , there are in totalp ( p + 1) / 2 different samples since we consider i + j and j + i to be the same sample.
A dataset Dis a set of non-repeating data samples.
We denote the full dataset as D 0 and split it into a trainingdataset D and a validation dataset D ′ , i.e., D (cid:83)
D ′ =
D 0 , D (cid:84)
D ′ = ∅ .
We define training datafraction =
| D | / | D 0 | where | · | denotes the cardinality of the set.
RQI: 0.0 — Accuracy - train: 1.0, validation: 0.
[Figure]
RQI: 0.6 — Accuracy - train: 1.0, validation: 0.
[Figure]
Generalization in toy addition
[Figure]
[Figure]
validation: 0.
[Figure]
Accuracy - train: 1.0, validation: 0.
[Figure]
Why Generalization: Representations and Dynamics
We can see that generalization appears to be linked to the emergence of highly-structured embeddingsin Figure 2.
In particular, Figure 2 (b) shows parallelograms in toy addition, and (d) shows a circlein toy modular addition.
We now restrict ourselves to the toy addition setup and formalize a notionof representation quality and show that it predicts the model’s performance.
We then develop aphysics-inspired effective theory of learning which can accurately predict the critical training set sizeand training trajectories of representations.
In physics, effective theories aim to describe observedphenomena without implying that the mechanisms used in the description necessarily have an exactcounterpart in the true underlying causes.
In our effective theory, we will model the dynamics ofrepresentation learning not as gradient descent of the true task loss but rather a simpler effective lossfunction ℓ eff which depends only on the representations and not on the decoder.
3.1 Representation quality predicts generalization for the toy model
A rigorous definition for structure in the learned representation is necessary.
We propose the followingdefinition,
Definition 1.
( i, j, m, n ) is a δ - parallelogram in the representation R ≡ [ E 0 , · · · , E p − 1 ] if
[Equation]
In the following derivations, we can take δ , which is a small threshold to tolerate numerical errors, tobe zero.
As long as training loss is effectively zero, any parallelogram ( i, j, m, n ) should satisfy i + j = m + n .
Suppose that this is not the case, i.e., suppose E i + E j = E
m
+ E n
but i + j ̸
= m + n , then
Y i + j =
Dec( E i + E j )
=
Dec( E m + E n ) =
Y m + n where the first and last equalities come fromthe zero training loss assumption.
However, since i + j ̸ = m + n , we have Y i + j ̸ = Y n + m (almost 87 surely in the regression task), a contradiction.
0 training data fraction 0 .0 .0 .0 .0 .1 .
A cc
M e m o r
i z a t i o n (a) 0 .0 .0 .
0 training set fraction 0 .0 .0 .0 .0 .1 .d
A cc (b) 0 .0 .0 .0 .0 .1 .d
Acc 0 .0 .0 .0 .0 .1 .
A cc (c)
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 to define the permissible parallelogram set associated with a training dataset D(“permissible” means consistent with 100% training accuracy) as
[Equation]
For simplicity, we denote P 0 ≡ P
( D 0 ) .
Given a representation R , we can check how manypermissible parallelograms actually exist in R within error δ , so we define the parallelogram setcorresponding to R as
[Equation]
For brevity we will write P ( R ) , suppressing the dependence on δ .
We define the representationquality index (RQI) as
[Equation]
We will use the term linear representation or linear structure to refer to a representation whoseembeddings are of the form E k
= a + k b ( k = 0 , · · · , p − 1; a , b ∈ R d in ) .
A linear representation 97 has RQI = 1 , while a random representation (sampled from, say, a normal dstribution) has RQI =
0with high probability.
Quantitatively, we denote the “predicted accuracy” (cid:100)
Acc as the accuracy achievable on the wholedataset given the representation R (see Appendix B for the full details).
In Figure 3, we see that thepredicted (cid:100)
Acc aligns well with the true accuracy Acc , establishing good evidence that representationleads to generalization.
We use an example to illustrate the origin of generalization here.
In the setupof Figure 2 (a), suppose the decoder can achieve zero training loss and E 6 + E 8 is a training samplehence Dec( E 6 + E 8 ) =
Y 14 .
At validation time, the decoder is tasked with predicting a validationsample E 5 + E 9 .
Since (5 , 9 , 6 , 8) forms a parallelogram such that E 5 + E 9 = E 6 + E 8 , the decodercan predict the validation sample correctly because Dec( E 5 + E 9 ) =
Dec( E 6 + E 8 ) =
Y 14 .
3.2 The critical training set size
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 e., Dec ∗ ( x 1 ) ̸
=
Dec ∗ ( x 2 ) for any x 1 ̸ = x 2 .
Then Proposition 1 provides a mechanism for the formation of parallelograms. Proposition 1.
If a training set D contains two samples ( i, j ) and ( m, n ) with i + j = m + n ,then M ∗ learns a representation R ∗ such that E i + E j = E m + E n , i.e., ( i, j, m, n
[Figure]
ili t y ( li n e a r s tr u c t u r e )
[Figure]
[Figure]
[Figure]
creditably
Proof.
Due to the zero training loss assumption, we have Dec ∗ ( E i + E j ) =
Y i + j = Y m + n =
Dec ∗ ( E m + E n ) .
Then the injectivity of Dec ∗ implies E i + [Equation]
The dynamics of the trained embedding parameters are determined by various factors interacting incomplex 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 thatthe dynamics of normalized quantities, namely, the normalized embeddings at time t , defined as˜ E ( t )
k
= E ( t ) k − µ t σ t , where µ t = 1 p (cid:80) k E ( t ) k and σ t = 1 p (cid:80) k | E ( t ) k − µ t | 2 , can be qualitatively describedby a simple effective loss.
We will assume that the normalized embedding vectors obey a gradientflow for an effective loss function of the form
[Equation]
[Equation]
5) where | · | denotes Euclidean vector norm.
Note that the embeddings do not collapse to the trivial 126 solution E 0 = · · · = E p − 1 = 0 unless initialized as such, because two conserved quantities exist, as 127 proven in Appendix D:
[Equation]
We shall now use the effective dynamics to explain empirical observations such as the existence of acritical training set size for generalization.
Degeneracy of ground states (loss optima)
We define ground states as those representationssatisfying ℓ eff = 0 , which requires the following linear equations to hold:
[Equation]
Since each embedding dimension obeys the same set of linear equations, we will assume, without lossof generality, that d in = 1 .
The dimension of the null space of A ( P ) , denoted as n 0 , is the number ofdegrees of freedom of the ground states.
Given a set of parallelograms implied by a training dataset
D , the nullity of A ( P ( D )) could be obtained by computing the singular values 0 ≤ σ 1 ≤ · · · ≤ σ p .
136 We always have n 0 ≥ 2 , i.e., σ 1 = σ 2 = 0 because the nullity of A ( P 0 ) , the set of linear equationsgiven by all possible parallelograms, is Nullity( A ( P 0 ))
= 2 which can be attributed to two degreesof freedom (translation and scaling).
If n 0 = 2 , the representation is unique up to translations andscaling factors, and the embeddings have the form E k
= a + k b .
Otherwise, when n 0 > 2 , therepresentation is not constrained enough such that all the embeddings lie on a line.
We present theoretical predictions alongside empirical results for addition ( p = 10 ) in Figure 4.
Asshown 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 datafraction and has a phase transition around r c = 0 .
Empirical results from training different modelsare shown in Figure 4(b).
The number of steps to reach RQI > 0 .
95 is seen to have a phase transitionat r c = 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 ℓ 0 as
[Equation]
[Equation]
Note that ℓ 0 = 12 R T HR , R =
[ E 0 , E 1 , · · · , E p − 1 ] , so the gradient descent is linear, i.e.,
[Equation]
If H has eigenvalues λ i =
σ 2 i (sorted in increasing order) and eigenvectors ¯ v i , and we have the initialcondition R ( t = 0)
= a i ¯ v i , then we have R ( t ) =
a i ¯ v i e − λ i t .
The first two eigenvalues vanish andt h = 1 /λ 3 determines the timescale for the slowest component to decrease by a factor of e .
Whenthe step size is η , the corresponding number of steps is n h
= t
h /η
= 1 / ( λ 3 η ) .
We verify the above analysis with empirical results.
Figure 4 (c)(d) show the trajectories obtainedfrom the effective theory and from neural network training, respectively.
The 1D neural representationin Figure 4(d) are manually normalized to zero mean and unit variance.
The two trajectories agreequalitatively, and it takes about 3 n h steps for two trajectories to converge to the linear structure.
Thequantitative differences might be due to the absence of the decoder in the effective theory, whichassumes the decoder to be adiabatic in the sense that it takes infinitesimal step sizes.
Limitations of the effective theory While our theory defines an effective loss based on the Euclideandistance between embeddings E i + E j and E n + E m , one could imagine generalizing the theory todefine a broader notion of parallogram given by some other metric on the representation space.
Forinstance, if we have a decoder like in Figure 2(c) then the distance between distinct representationswithin the same “pizza slice” is low, meaning that representations arranged not in parallelogramsw.r.t.
the Euclidean metric may be parallelograms with respect to the metric defined by the decoder.
Delayed Generalization: A Phase Diagram
So far, we have (1) observed empirically that generalization on algorithmic datasets corresponds withthe emergence of well-structured representations, (2) defined a notion of representation quality in atoy setting and shown that it predicts generalization, and (3) developed an effective theory to describethe learning dynamics of the representations in the same toy setting.
We now study how optimizerhyperparameters affect high-level learning performance.
In particular, we develop phase diagrams forhow learning performance depend on the representation learning rate, decoder learning rate and thedecoder weight decay.
These parameters are of interest since they most explicitly regulate a kind ofcompetition between the encoder and decoder, as we elaborate below.
4.1 Phase diagram of a toy model
Training details We update the representation and the decoder with different optimizers.
For the1D embeddings, we use the Adam optimizer with learning rate [10 − 4 , 10 − 2 ] and zero weight decay.
For the decoder, we use an AdamW optimizer with the learning rate in [10 − 4 , 10 − 2 ] and the weightdecay in [0 , 10] (regression) or [0 , 20] (classification).
For training/validation spliting, we choose45/10 for non-modulo addition ( p = 10 ) and 24/12 for the permutation group S 3 .
We hard-codeaddition or matrix multiplication (details in Appendix F) in the decoder for the addition group andthe permutation group, respectively.
For each choice of learning rate and weight decay, we compute the number of steps to reach high(90%) training/validation accuracy The 2D plane is split into four phases: comprehension , grokking ,memorization and confusion , defined in Table 1.
Both comprehension and grokking are able togeneralize (in the “Goldilocks zone"), although the grokking phase has delayed generalization.
Memorization is also called overfitting, and confusion means failure to even memorize training data.
Figure 5 shows the phase diagrams for the addition group and the permutation group.
They displayquite rich phenomenon.
1e1e- 1e1e-2 decoder learning rate 1e1e- 1e1e- r e p r e s e n t a t i o n l e a r n i n g r a t e c o m p r e h e n s i o n memorization grokking (a) Addition, regression 1e1e- 1e1e-2 decoder learning rate   w e i g h t d ec a y c o m p r
e
h e n s i o n memorization grokking confusion (b) Addition, regression 1e1e- 1e1e-2 decoder learning rate
w e
i g h t d ec a y comprehension m e m o r
i z a t
i o n grokking c o n f u s i o n (c) Addition, classification 1e1e- 1e1e-2 decoder learning rate
w e i g h t d ec a y comprehension memorization
g r o kk i n g confusion
[Figure]
[Table]
within
[Table]
step(validation acc> [Table]
[Table]
<
[Table]
Comprehension
Yes Yes
Yes
Grokking
Yes Yes
No
Memorization
Yes
No
Not Applicable
Confusion
No No
Not Applicable
Competition between representation learning and decoder overfitting In the regression setup ofthe addition dataset, we show how the competition between representation learning rate and decoderlearning rate lead to different learning phases in Figure 5 (a).
As expected, a fast decoder coupledwith slow representation learning (bottom right) lead to memorization.
In the opposite extreme,although an extremely slow decoder coupled with fast representation learning (top left) will generalizein 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.
Drawing from an analogy to physical systems, one can think of embedding vectors as a group ofparticles.
In our effective theory from Section 3.2, internal particle interactions promote the formationof structure, while external forces inhibit it.
The decoder plays the role of an environment exertingexternal forces on the embeddings.
If the decoder changes slowly/fast, the magnitude of the externalforces are small/large, leading to better/worse representations.
Universality of phase diagrams We fix representation learning to be 10 − 3 and sweep insteaddecoder weight decay in Figure 5 (b)(c)(d).
The phase diagrams correspond to addition regression(b), addition classification (c) and permutation regression (d), respectively.
Common phenomenonemerge from these different tasks: (i) they all include four phases; (ii) The top right corner (a fast andcapable decoder) is the memorization phase; (iii) the bottom right corner (a fast and simple decoder) isthe confusion phase; (iv) grokking is sandwiched between comprehension and memorization, whichseems to imply that it is an undesirable phase that stems from improperly tuned hyperparameters.
Due to limited space, we include phase diagrams of other hyperparameters (batch size, intializationscale and weight decay on representation) in Appendix E.
4.2 Beyond the toy model
We conjecture that many of the principles which we saw dictate the training dynamics in the toymodel also apply more generally.
Below, we will see how our framework generalizes to transformerarchitectures for the task of addition modulo p , a minimal reproducible example of the originalgrokking paper [1].
We first encode p = 53 integers into 256D learnable embeddings, then pass two integers to a decoderonly transformer architecture.
For simplicity, we do not encode the operation symbols here.
Theoutputs from the last layer are concatenated and passed to a linear layer for classification.
Trainingboth the encoder and the decoder with the same optimizer (i.e., with the same hyperparameters)leads to the grokking phenomenon.
Generalization appears much earlier once we lower the effectivedecoder capacity with weight decay (full phase diagram in Figure 6).
[Figure]
o c h s t o% a cc u r a c y validation train validation - train 3e1e- 4e-2 learning rate 1e2e- 2e+ 3e+ w e i g h t d e c a y c o m p r e h e n s i o n memorization grokking confusion
Early on, the model is able to perfectly fit the training set while having no generalization.
We studythe 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 perfectgeneralization coincides with the PCA projecting the embeddings into a circle in 2D. Of course, nochoice of dimensionality reduction is guaranteed to find any structure, and thus, it is challengingto show explicitly that generalization only occurs when a structure exists.
Nevertheless, the factthat, when coupled with the implicit regularization of the optimizer for sparse solutions, such a clearstructure appears in a simple PCA so quickly at generalization time suggests that our analysis inthe toy setting is applicable here as well.
This is also seen in the evolution of the entropy of theexplained variance ratio in the PCA of the embeddings (defined as S = − (cid:80) i σ i log σ i where σ i is 231 the fractional variance explained by the i th principal component).
As seen in Figure 6, the entropyincreases up to generalization time then decreases drastically afterwards which would be consistentwith the conjecture that generalization occurs when a low-dimensional structure is discovered.
Thedecoder then primarily relies on the information in this low-dimensional manifold and essentially“prunes” the rest of the high-dimensional embedding space.
In Figure 6 (right), we show a comparable phase diagram to Figure 5 evaluated now in the transformersetting.
Note that, as opposed to the setting in [1], weight decay has only been applied to the decoderand not to the embedding layer.
Contrary to the toy model, a certain amount of weight decay proves beneficial to generalization and speeds it up significantly.
We conjecture that this difference comesfrom the different embedding dimensions.
With a highly over-parameterized setting, a non-zeroweight decay gives a crucial incentive to reduce complexity in the decoder and help generalize infewer steps.
This is subject to further investigation.
We also explore the effect of dropout layers inthe decoder blocks of the transformer.
With a significant dropout rate, the generalization time can bebrought down to under 10 3 steps and the grokking phenomenon vanishes completely.
The overalltrend suggests that constraining the decoder with the same tools used to avoid overfitting reducesgeneralization time and can avoid the grokking phenomenon.
Related work
We are not aware of other formal attempts to understand grokking, though some authors have providedspeculative, informal accounts.
For instance, [2] describes training as a random walk over parameters,and draws connections to Solomonoff induction.
Separately, [3] suggest that simpler functions whichgeneralize achieve lower loss than complicated functions which memorize, providing a training signalencouraging generalization.
More broadly, our work is related to the following research directions:
Learning mathematical structures [4] trains a neural network to learn arithmetic operation frompictures of digits, but they do not observe grokking due to their abundant training data.
Beyondarithmetic relations, machine learning has been applied to learn other mathematical structures,including geometry [5], knot theory [6] and group theory [7].
Double descent Grokking is somewhat reminiscent of the phenomena of “epoch-wise” doubledescent [8], where generalization can improve after a period of overfitting.
[9] find that regularizationcan mitigate double descent, similar perhaps to how weight decay influences grokking.
Representation learning Representation learning lies at the core of machine learning [10–13].
Representation quality is usually measured by (perhaps vague) semantic meanings or performance ondownstream tasks.
In our study, the simplicity of arithmetic datasets allows us to define representationquality and study evolution of representations in a quantitative way.
Physics of learning Physics-inspired tools have been proved useful to understand machine learningfrom a theoretical perspective.
These tools include effective theories [14, 15], conservation laws [16]and free energy principle [17].
However, none of these works aim to develop an effective theory ofrepresentation learning.
We have argued that grokking occurs when models learn a particularly structured representationof their inputs.
This connects a low-level understanding of the internal operation of our modelswith a high-level phenomenon: an increase in validation accuracy.
In a recent work, researchers at Anthropic [18], connect a sudden decrease in loss during training with the emergence of inductionheads within their models.
They analogize their work to statistical physics , since it bridges a“microscopic”, mechanistic understanding of networks with “macroscopic” facts about overall modelperformance.
We view our work in the same vein, but we make the additional contribution that wedevelop, in a toy setting, an analytic model of low-level representation learning dynamics .
Conclusion
We have shown how, in both toy models and more general settings, that representation enablesgeneralization when it reflects structure in the data.
We developed an effective theory of representationlearning dynamics (in a toy setting) which predicts the critical dependence of learning on the trainingdata fraction.
We then presented phase transitions between comprehension, grokking, memorizationand confusion where the learning “phase” depends on the decoder capacity and learning rate indecoder-only architectures.
While we have mostly focused on a toy model, we find preliminaryevidence that our results generalize to the setting of [1].
Our work can be viewed as a step towards a statistical physics of deep learning , a theory whichconnects the “microphysics” of low-level network dynamics with the “thermodynamics” of highlevel model behavior.
We view the application of theoretical tools from physics, such as effectivetheories [19], to be a rich area for further work.
The broader impact of such work, if successful,would be to make models more transparent and predictable [18, 20, 21], crucial to the task of ensuringthe safety of advanced AI systems. Checklist For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes] We lay out our “microscopic” analysis and effectivetheory in Section 3. For a toy model, we establish the connection between wellstructured representations and generalization in 3.1 and describe an effective theory of the learning dynamics of these representations
in 3.2, which we find predicts the
critical training set size.
We perform our “macroscopic” analysis of phase diagrams in
Section 4.
(b) Did you describe the limitations of your work?
[Yes] Limitations are described through-
out the text where applicable, and also explicitly in the Appendix.
(c) Did you discuss any potential negative societal impacts of your work?
[No] Our work
is theoretical.
While some of our insights may eventually lead to more capable models,
which could be misused, we don’t think our work introduces any specific risks for
misuse.
On the other hand, work like ours trying to gain understanding of deep learning
phenomena could prove useful for AI safety.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] The theory we
develop in this paper is, in the language of physics, an effective theory.
It is simpler
than the full setup where grokking was initially observed, but this simplicity allows us
to develop a precise mathematical description of representation quality and learning
dynamics.
We make clear that our theoretical analysis is specific to our toy model.
(b) Did you include complete proofs of all theoretical results?
[Yes] We have relatively
few theoretical results, contained within our discussion of our effective theory on the
toy model 3.
We provided short proofs of those results.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
imental results (either in the supplemental material or as a URL)?
[Yes]
In the final
(non-anonymized) version of the paper we will include a working link to a public
repository containing code for reproducing our results and figures.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
(c) Did you report error bars (e.g., with respect to the random seed after running ex-
periments multiple times)?
[No] It is difficult to visualize uncertainty for our phase
diagrams.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[No] Our experiments can be reproduced
on a single GPU within a week.
Since this is not a substantial amount of compute, andour study is about gaining scientific understanding rather than achieving SOTA, we did
not deem it important to report this in the text.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] Some of our code
was based off of code released by the authors of the original paper on grokking [1].
(b) Did you mention the license of the assets?
[No] The code for reproducing our experi-
ments will be made available on GitHub under the MIT license.
We release no otherassets.(c)
Did you include any new assets either in the supplemental material or as a URL?
[Yes]In the full, non-anonymized publication, a working link will be included to the codefor reproducing our experiments.
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or 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]
[Table]
[Table]
[Table]
[Table]
[Table]
N
[Table] Appendix A Applicability of our toy setting In the main paper, we focused on the toy setting with (1) the addition dataset and (2) the additionoperation hard coded in the decoder.
Although both simplifications appear to have quite limitedapplicability, we argue below that the analysis of the toy setting can actually apply to all Abeliangroups.
The addition dataset is the building block of all Abelian groups A cyclic group is a group thatis generated by a single element.
A finite cyclic group with order n is C n = { e, g, g 2 , · · · , g n − 1 } 420 where e is the identify element and g is the generator and g i =
g j whenever i =
j (mod n ) .
Themodulo addition and { 0 , 1 , · · · , n − 1 } form a cyclic group with e = 0 and g can be any number q 422 coprime to n such that ( q, n ) = 1 .
Since algorithmic datasets contain only symbolic but no arithmeticinformation, the datasets of modulo addition could apply to all other cyclic groups, e.g., modulomultiplication and discrete rotation groups in 2D.
Although not all Abelian groups are cyclic, a finite Abelian group G can be always decomposed intoa direct product of k cyclic groups G = C n 1 × C n 2 · · · C n k .
So after training k neural networks with 427 each handling one cyclic group separately, it is easy to construct a larger neural network that handlesthe whole Abelian group.
The addition operation is valid for all Abelian groups It is proved in [22] that for a permutationinvariant function f ( x 1 , x 2 , · · · , x n ) , there exists ρ and ϕ such thatf ( x 1 , x 2 , · · · , x n )
=
ρ [ n (cid:88) i = ϕ
( x i )] , (10) or f ( x 1 , x 2 ) =
ρ ( ϕ ( x 1 ) + ϕ ( x 2 )) for n = 2 .
Notice that ϕ
( x i ) corresponds to the embedding vector
E i , ρ corresponds to the decoder.
The addition operator naturally emerges from the commutativityof the operator, not restricting the operator itself to be addition.
For example, multiplication oftwo numbers x
and x 2 can be written as x 1 x 2 = exp(ln( x 1 ) + ln( x 2 )) where ρ ( x ) =
exp( x ) andϕ ( x ) = ln( x ) .
B Definition of (cid:100) Acc
Given a training set D and a representation R , if ( i, j ) is a validation sample, can the neural networkcorrectly predict its output, i.e., Dec( E i + E j )
=
Y i + j ?
Since neural network has never seen ( i, j )in the training set, one possible mechanism of induction is through
[Equation]
The first equality Dec( E i + E j ) =
Dec( E m + E n ) holds only when E i + E j = E m + E n (i.e.,( i, j, m, n ) is a parallelogram).
The second equality Dec( E m + E n ) =
Y m+n , holds when ( m, n ) intraining set, i.e., ( m, n ) ∈ D , under the zero training loss assumption.
Rigorously, given a training set D and a parallelogram set P (which can be calculated from R ), we collect all zero loss samplesin an augmented training set D [Equation]
[Equation]
Keeping D fixed, a larger P would probably produce a larger D , i.e., if P 1 ⊆ P 2 , then D ( D, P 1 ) ⊆
D ( P, P 2 ) , which is why in Eq.
(3) our defined RQI ∝ | P | gets its name “representation quality 447 index", because higher RQI normally means better generalization.
Finally, the expected accuracyfrom a dataset D and a parallelogram set P is:
[Equation]
which is the estimated accuracy (of the full dataset), and P = P ( R ) is defined on the representationafter training.
On the other hand, accuracy Acc can be accessed empirically from trained neuralnetwork.
We verified Acc ≈ (cid:100)
Acc in a toy setup (addition dataset p = 10 , 1D embedding space, hardcode addition), as shown in FIG.
3 (c).
FIG.
3 (a)(b) show Acc and (cid:100)
Acc as a function of training setratio, with each dot corresponding to a different random seed.
The dashed red diagonal correspondsto memorization of the training set, and the vertical gap refers to generalization.
Although the agreement is good for 1D embedding space, we do not expect such agreement cantrivially extend to high dimensional embedding space.
In high dimensions, our definition of RQI istoo restrictive.
For example, suppose we have an embedding space with N dimensions.
Although therepresentation may form a linear structure in the first dimension, the representation can be arbitraryin other N − 1 dimensions, leading to RQI ≈ 0 .
However, the model may still generalize well if thedecoder learns to keep only the useful information (the 1st dimension) and drop all other uselessinformation (other 31 dimensions).
It would be interesting to investigate how to define an RQI thattakes into account the role of decoder in future works.
C The gap of a realistic model M and the ideal model M ∗
Realistic models M usually form fewer number of parallelograms than ideal models M ∗ .
In thissection, we analyze the properties of ideal models and calculated ideal RQI and ideal accuracy,which set upper bounds for empirical RQI and accuracy.
The upper bound relations are verified vianumerical experiments in Fig. 8.
Similar to Eq.
(12) where some validation samples can be derived from training samples, wedemonstrate how implicit parallelograms can be ‘derived’ from explicit ones in P
( D ) .
The so-calledderivation follows a simple geometric argument that: if A 1 B 1 is equal and parallel to A 2 B 2 , and
A 2 B 2 is equal and parallel to A 3 B 3 , then we can deduce that A 1 B 1 is equal and parallel to A 3 B 3(hence ( A 1 , B 2 , A 2 , B 1 ) is a parallelogram).
Recall that a parallelogram ( i, j, m, n ) is equivalent to E i + E j = E m + E n ( ∗ ) .
So we areequivalently asking if equation ( ∗ ) can be expressed as a linear combination of equations in
A ( P 0 ( D )) .
If yes, then ( ∗ ) is dependent on A ( P 0 ( D ))
(defined in Eq.
(7)), i.e., A ( P 0 ( D )) and
A ( P 0 ( D ) (cid:83) ( i, j, m, n )) should have the same rank.
We augment P
( D ) by adding implicit parallelograms, and denote the augmented parallelogram set as
P ( D ) =
P {+
+}
( D
[Equation]
(cid:91)
[Equation]
We need to emphasize that an assumption behind Eq.
(14) is that we have an ideal model M ∗ .
Whenthe model is not ideal, e.g., when the injectivity of the encoder breaks down, fewer parallelogramsare expected to form, i.e.,
[Equation]
The inequality is saying, whenever a parallelogram is formed in the representation after training, thereason is hidden in the training set.
This is not a strict argument, but rather a belief that today’s neuralnetworks can only copy what datasets (explicitly or implicitly) tell it to do, without any autonomouscreativity or intelligence.
For simplicity we call this belief Alexander Principle .
In very rare caseswhen something lucky happens (e.g., neural networks are initialized at approximate correct weights),
Alexander principle may be violated.
Alexander principle sets an upper bound for RQI :
[Equation]
0 training data fraction 0 .0 .0 .0 .0 .1 .
R Q I (a) 0 .0 .0 .0 .0 .1 .
0 training data fraction 0 .0 .0 .0 .0 .1 .
R Q I (b) 0 .0 .0 .
RQI 0 .0 .0 .0 .0 .1 .
R Q
I
I d e a l A l go r
i t h m
A l e x a nd e r
P r i n c i p l e (c)
[Figure]
[Figure] [Figure]
fraction
[Figure]
A cc
M e m o r
i z a t
i o n (d) 0 .0 .0 .0 .0 .1 .
0 training data fraction 0 .0 .0 .0 .0 .1 .
A cc
M e m o r
i z a t i o n (e) 0 .0 .0 .
Acc 0 .0 .0 .0 .0 .1 .
A cc
I d e a l A l go r
i t h m
A l e x a nd e r
P r i n c i p l e (f) c)(f), RQI and Acc of the ideal algorithm sets upper bounds for those of the realistic algorithm.
and sets an upper bound for
[Equation]
In Figure 8 (c)(f), we verify Eq.
(16) and Eq.
We choose δ = 0 .
01 to compute RQI( R , δ ).
Wefind the trained models are usually far from being ideal, although we already include a few usefultricks proposed in Section 4 to enhance representation learning.
It would be an interesting futuredirection to develop better algorithms so that the gap due to Alexander principle can be reducedor even closed.
In Figure 8 (a)(b)(d)(e), four quantities (RQI, RQI , Acc, Acc ) as functions of thetraining data fraction are shown, each dot corresponding to one random seed.
It is interesting tonote that it is possible to have RQI = 1 only with < 40% training data, i.e., 55 × 0 .
4 = 22 samples,agreeing with our observation in Section 3.
Realistic representations Suppose an ideal model M ∗ and a realistic model M which train on thetraining set D give the representation R ∗ and R , respectively.
What is the relationship between Rand R ∗ ?
Due to the Alexander principle we know P ( R ) ⊆
P ( D ) =
P ( R ∗ ) .
This means R ∗ hasmore parallelograms than R , hence R ∗ has fewer degrees of freedom than R .
We illustrate with the toy case p = 4 .
The whole dataset contains p ( p + 1) / 2 = 10 samples, i.e.,
D
= { (0 , 0) , (0 , 1) , (0 , 2) , (0 , 3) , (1 , 1) , (1 , 2) , (1 , 3) , (2 , 2) , (2 , 3) , (3 , 3) } .
The parallelogram set contains only three elements, i.e.,
[Equation]
Or equivalently the equation set
[Equation]
Pictorially, we can split all possible subsets { A | A ⊆
A 0 } into different levels, each level defined by| A | (the number of elements).
A subset A 1 in the i th level points an direct arrow to another subset
A 2 in the ( i + 1) th level if A 2 ⊂ A 1 , and we say A 2 is a child of A 1 , and A 1 is a parent of A 2 .
Each subset A can determine a representation R with n ( A ) degrees of freedom.
So R should be adescendant of R ∗ , and n ( R ∗ ) ≤ n ( R ) .
Numerically, n ( A ) is equal to the dimension of the null spaceof A .
Suppose we have a training set
D = { (0 , 2) , (1 , 1) , (0 , 3) , (1 , 2) , (1 , 3) , (2 , 2) } , (21) and correspondingly P ( D ) =
P 0 , A ( P ) =
A 0 .
So an ideal model M ∗ will have the linear structure
E k
=
a + k b (see Figure 9 leftmost).
However, a realistic model M may produce any descendantsof the linear structure, depending on various hyperparameters and even random seeds.
In Figure 10, we show our algorithms actually generates all possible representations.
We havetwo settings: (1) fast decoder ( η 1 , η 2 ) =
(10 − 3 , 10 − 2 ) (Figure 10 left), and (2) relatively slowdecoder ( η 1 , η 2 ) =
(10 − 2 , 10 − 3 ) (Figure 10) right).
The relatively slow decoder produces betterrepresentations (in the sense of higher RQI) than a fast decoder, agreeing with our observation in
Section 4.
D Conservation laws of the effective theory
Recall that the effective loss function
[Equation]
where ℓ 0 and Z 0 are both quadratic functions of R = { E 0 , · · · , E p − 1 } , and ℓ eff = 0 remains zero 521 under rescaling and translation E ′ i = a E i + b .
The representation vector E i evolves according to thegradient descent
[Equation]
We will prove the following two quantities are conserved:
[Equation]
Eq.
(22) and Eq.
(23) give
[Equation]
Then
[Equation]
where the last equation uses the fact that
[Equation]
The conservation of Z 0 prohibits the representation from collapsing to zero.
Now that we havedemonstrated that Z 0 is a conserved quantity, we can also show
[Equation]
The last equality holds because the two summations can be swapped and
[Equation]
[Equation]
E More phase diagrams of the toy setup
We study another three hyperparameters in the toy setup by showing phase diagrams similar to
Figure 5.
The toy setup is: (1) addition without modulo ( p = 10 ); (2) training/validation is split into45/10; (3) hard code addition; (4) 1D embedding.
In the following experiments, the decoder is an
1e1e- 1e-2 learning rate   b a t c h
s i ze c
o m p r
e
h e n s i o n memorization grokking confusion
Addition group (regression) 1e1e- 1e-2 learning rate   b a t c h
s i ze c
o m p r
e
h e n s i o n memorization g r
o kk i n g confusion
Addition group (classiﬁcation)
MLP with size 1-200-200-30.
The representation and the encoder are optimized with AdamW withdifferent hyperparameters.
The learning rate of the representation is 10 − 3 .
We sweep the learningrate of the decoder in range [10 − 4 , 10 − 2 ] as the x axis, and sweep another hyperparameter as they axis.
By default, we use full batch size 45, initialization scale s = 1 and zero weight decay ofrepresentation.
Batch size controls the amount of noise in the training dynamics.
In Figure 11, the grokking regionappears at the top left of the phase diagram (small decoder learning rate and small batch size).
However, large batch size (with small learning rate) leads to comprehension, implying that smallerbatch size seems harmful.
This makes sense since to get crystals (good structures) in experiments,one needs a freezer which gradually decreases temperature, rather than something perturbing thesystem with noise.
Initialization scale controls distances among embedding vectors at initialization.
We initializecomponents of embedding vectors from independent uniform distribution U [ − s/ 2 , s/ 2] where sis called the initialization scale.
Shown in Figure 12, it is beneficial to use a smaller initializationscale.
This agrees with the physical intuition that closer particles are more likely to interact and formstructures.
For example, the distances among molecules in ice are much smaller than distances in gas.
Representation weight decay controls the magnitude of embedding vectors.
Shown in Figure 13,we see the representation weight decay in general does not affect model performance much.
F General groupsF.1 Theory
We focused on Abelian groups for the most part of the paper.
This is, however, simply due topedagogical reasons.
In this section, we show that it is straight-forward to extend definitions ofparallelograms and representation quality index (RQI) to general non-Abelian groups.
We will alsoshow that most (if not all) qualitative results for the addition group also apply to the permutationgroup.
Matrix representation for general groups Let us first review the definition of group representation.
A representation of a group G on a vector space V is a group homomorphism from G to GL( V ) , thegeneral linear group on V .
That is, a representation is a map ρ : G → GL( V ) such that
[Equation]
[Figure] [Figure]
i n i t i a li z a t i o n s c a l e comprehension memorization confusion
Addition group (regression)
[Figure]
[Figure]
i n i t i a li z a t i o n s c a l e comprehension
[Figure]
confusion
Addition group (classiﬁcation)
[Figure]
w e i g h t d ec a y (r e p r e s e n t a t i o n ) c
o m p
r
e h e n s i o n
[Figure]
g r
o
kk
i n g
Addition group
[Figure]
regression)
[Figure]
n
[Figure]
) c
[Figure]
m p r
[Figure]
h
[Figure]
n m
[Figure]
m
[Figure]
g r
o
kk
i n g c o n f u s i o n
[Figure]
In the case V is of finite dimension n , it is common to identify GL( V ) with n by n invertible matrices.
The punchline is that: each group element can be represented as a matrix, and the binary operation isrepresented as matrix multiplication.
A new architecture for general groups Inspired by the matrix representation, we embed eachgroup element a as a learnable matrix E a ∈ R d × d (as opposed to a vector), and manually do matrixmultiplication before sending the product to the deocder for regression or classification.
Moreconcretly, for a ◦ b = c , our architecture takes as input two embedding matrices E a and E b andaims to predict Y c such that Y c
=
Dec( E a E b ) , where E a E b means the matrix multiplication of E aand E b .
The goal of this simplication is to disentangle learning the representation and learning thearithmetic operation (i.e, the matrix multiplication).
We will show that, even with this simplification,we are still able to reproduce the characteristic grokking behavior and other rich phenomenon.
Generalized parallelograms we define generalized parallelograms: ( a, b, c, d ) is a generalizedparallelogram in the representation if || E a E b −
E c E d || 2 F ≤ δ , where δ > 0 is a threshold to tolerate 577 numerical errors.
Before presenting the numerical results for the permutation group, we show anintuitive picture about how new parallelograms can be deduced from old ones for general groups,which is the key to generalization.
Deduction of parallelograms
We first recall the case of the Abelian group (e.g., addition group).
As shown in Figure 14, when( a, d, b, c ) and ( c, f, d, e ) are two parallelograms, we have
[Equation]
We can derive that E a + E f = E b
+ E e implying that ( a, f, b, e ) is also a parallelogram.
That is, for
Abelian groups, two parallelograms are needed to deduce a new parallelogram.
For the non-Abelian group, if we have only two parallelograms such that
[Equation]
) we have E − 1 b E a = E c E − 1 d = E − 1 f E e , but this does not lead to something like E f E a = E e E b ,hence useless for generalization.
However, if we have a third parallelogram such that
[Equation]
31) we have E − 1 b E a = E c E − 1 d = E − 1 f
E
e
=
E g E − 1 h , equivalent to E a E h
= E b E g , thus establishing anew parallelogram ( a, h, b, g ) .
That is, for non-Abelian groups, three parallelograms are needed todeduce a new parallelogram.
F.2 Numerical Results
In this section, we conduct numerical experiments on a simple non-abelian group: the permutationgroup S 3 .
The group has 6 group elements, hence the full dataset contains 36 samples.
We embed eachgroup element a into a learnable 3 × 3 embedding matrix E a .
We adopt the new architecture describedin the above subsection: we hard code matrix multiplication of two input embedding matrices beforefeeding to the decoder.
After defining the generalized parallelogram in the last subsection, we cancontinue to define RQI (as in Section 3) and predict accuracy (cid:100)
Acc from representation (as in appendix
B).
We also compute the number of steps needed to reach RQI = 0 .
Representation We flatten each embedding matrix into a vector, and apply principal componentanalysis (PCA) to the vectors.
We show the first three principal components of these group elementsin Figure 15.
On the plane of PC1 and PC3 , the six points are organized as a hexagon.
PC 1 3 2 1 0 1
PC 2 1.51.00.50.00.51.01.
PC1.51.00.50.00.
[0, 1, 2] [0, 2, 1] [1, 0, 2] [1, 2, 0] [2, 0, 1] [2, 1, 0]      3 PC1 1.
P C[0, 1, 2] [0, 2, 1] [1, 0, 2] [1, 2, 0] [2, 0, 1] [2, 1, 0]      3 PC1 1.
P C[0, 1, 2] [0, 2, 1] [1, 0, 2] [1, 2, 0] [2, 0, 1] [2, 1, 0] 1.5 1.0 0.5 0.0 0.5 1.0 1.5 PC2 1.
P C[0, 1, 2] [0, 2, 1] [1, 0, 2] [1, 2, 0] [2, 0, 1] [2, 1, 0] [Figure] [Figure]
[Figure]
s t e p
[Figure]
t
[Figure]
R Q I >
[Figure]
r c =
[Figure]
(b) 0 . [Figure]
0 training data fraction 0 .0 .
A cc (c) 0 .0 .0 .0 .1 .
0 training data fraction 0 .0 .
9 1 .d
A cc (d) 0 .0 .0 .0 .1 .d
Acc 0 .0 .
A cc (e)
RQI In Figure 16(a), we show RQI as a function of training data fraction.
For each training datafraction, we run 11 random seeds (shown as scatter points), and the blue
Steps to reach RQI = 0 .
In Figure 16(b), we whow the steps to reach RQI > 0 .
95 as a function oftraining data fraction, and find a phase transition at r = r c = 0 .
The blue line corresponds to thebest model (smallest number of steps).
Accuracy
The real accuracy Acc is shown in Figure 16(c), while the predicted accuracy (cid:100)
Acc(calculated from RQI)
is shown in Figure 16(d).
Their comparison is shown in (e): (cid:100)
Acc is a lowerbound of Acc , implying that there must be some generalization mechanism beyond RQI.
Phase diagram We investigate how the model performance varies under the change of two knobs:decoder learning rate and decoder weight decay.
We calculate the number of steps to training accuracy≥ 0 .
9 and validation accuracy ≥ 0 .
9 , respectively, shown in Figure 5 (d).
[-With these modifications to training set size and initialization scale, we train a depth-3 width-
-]
[-The adiabatic limit Using the standard initialization (e.g., Xavier initialization) of neural networks,we have A T 
-]
[-A 
-]
[-where P 
-]