Towards Understanding Grokking: An Effective Theory of Representation Learning
Anonymous Author(s)
AffiliationAddress email
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, raising three key puzzles:
Q
The origin of generalization : When trained on the algorithmic datasets where grokking occurs, how do models generalize at all?
Q
The critical training size : Why does the training time needed to “grok” (generalize) diverge asthe 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 input embeddings,
i.e. a representation that has the appropriate structure for the task and which can be predictedfrom the theory in Section 3.
See Figures 1 and 2.
[Figure] A2 The critical training set size corresponds to the least amount of training data that can determinesuch a representation (which, in some cases, is unique up to linear transformations).
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(mean-squared error) and classification (cross-entropy).
The basic setup is as follows: our modeltakes as input the symbols a, b and maps them to trainable embedding vectors E a , E b ∈ R d in .
Itthen sums E a , E b and sends the resulting vector through a “decoder” MLP.
The target output vector,denoted Y c ∈ R d out is a fixed random vector (regression task) or a one-hot vector (classificationtask).
Our model architecture can therefore be compactly described as ( a, b ) (cid:55)→
Dec( E a + E b ) ,where the embeddings E ∗ and the decoder are trainable.
Despite its simplicity, this toy model cangeneralize to all abelian groups (discussed in Appendix A).
In sections 3-4.1, we consider only thebinary operation of addition.
We consider modular addition in Section 4.2 to generalize some of ourresults to a transformer architecture and study 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]
train
[Figure]
1.0, 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.
The concept of an effective theory in physics is similar tomodel reduction in computational methods, aiming to describe complex phenomena with simple yetintuitive pictures.
In our effective theory, we will model the dynamics of representation learning notas gradient descent of the true task loss but rather a simpler effective loss function ℓ eff which dependsonly on the representations in embedding space 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 88 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 data
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 98 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 thatthe predicted (cid:100)
Acc aligns well with the true accuracy Acc , establishing good evidence that structuredrepresentation of input embeddings leads to generalization.
We use an example to illustrate the originof generalization here.
In the setup of Figure 2 (a), suppose the decoder can achieve zero trainingloss and E 6 + E 8 is a training sample hence Dec( E 6 + E 8 ) =
Y 14 .
At validation time, the decoderis tasked with predicting a validation sample E 5 + E 9 .
Since (5 , 9 , 6 , 8) forms a parallelogramsuch that E 5 + E 9 = E 6 + E 8 , the decoder can predict the validation sample correctly because Dec( E 5 + E 9 ) =
Dec( E 6 + E 8 ) =
Y 14 .
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.e., Dec ∗ ( x 1 ) ̸
=
Dec ∗ ( x 2 ) for any x 1 ̸ = x 2 .
Then Proposition 1 provides a mechanism for the formation of parallelograms.
[Figure]
ili t y ( li n e a r s tr u c t u r e )
[Figure]
[Figure]
1D creditably well. 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 ) forms aparallelogram.
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 + E j = E m + E n .
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
[Equation]
, 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 (in the physics effective theory sense).
We will assume that the normalizedembedding vectors obey a gradient flow for an effective loss function of the form
[Equation]
[Equation]
where | · | denotes Euclidean vector norm.
Note that the embeddings do not collapse to the trivial 129 solution E
[Equation]
[Equation] E [Equation] 1 = 0 unless initialized as such, because two conserved quantities exist, as 130 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 .
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 degrees of 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 aunique 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]
Note that ℓ 0 = 12 R T HR , R =
[ E 0 , E 1 , · · · , E p − 1 ] , so the gradient descent is linear, i.e., [Equation]
[Equation]
If H has eigenvalues λ i =
σ 2 i (sorted in increasing order) and eigenvectors ¯ v i , and we have the initialcondition R ( t = 0)
= (cid:80) i a i ¯ v i , then we have R ( t ) =
(cid:80) i a i ¯ v i e − λ i t .
The first two eigenvaluesvanish and t h = 1 /λ 3 determines the timescale for the slowest component to decrease by a factorof e .
When the step size is η , the corresponding number of steps is n h
= t
h /η
= 1 / ( λ 3 η ) . See
Appendix G for more details on the dependence of λ 3 on training set size.
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.
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 [Figure]
[Figure]
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(
[Table]
[Table]
Comprehension
Yes Yes
Yes
Grokking
Yes Yes
No
Memorization
Yes
No
Not Applicable
Confusion
No No
Not Applicable
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 phenomena.
Competition between representation learning and decoder overfitting In the regression setupof the addition dataset, we show how the competition between representation learning and decoderlearning (which depend on both learning rate and weight decay, among other things) lead to differentlearning phases in Figure 5 (a).
As expected, a fast decoder coupled with slow representation learning(bottom right) lead to memorization.
In the opposite extreme, although an extremely slow decodercoupled 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 representationlearning 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.1, the dynamics of the particles are described only bytheir relative positions, in that sense, structure forms mainly due to these interaction.
The decoderplays the role of an environment exerting external forces on the embeddings.
If the magnitude of theexternal forces are small/large one can expect 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.
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]
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 234 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.
The decoder then primarily relies on the information in this low-dimensional manifold and essentially“prunes” the rest of the high-dimensional embedding space. See Appendix J for more details.
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 provesbeneficial 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 layersin the decoder blocks of the transformer.
With a significant dropout rate, the generalization timecan be brought down to under 10 3 steps and the grokking phenomenon vanishes completely.
Theoverall trend suggests that constraining the decoder with the same tools used to avoid overfittingreduces generalization time and can avoid the grokking phenomenon. This is also observed in animage classification task where we were able to induce grokking. See Appendix I for more details.
Related work
We are not aware of other formal attempts to understand grokking, though some authors have providedspeculative, informal accounts [2, 3].
Our work is related to the following broad research directions:
Double descent Grokking is somewhat reminiscent of the phenomena of “epoch-wise” doubledescent [4], where generalization can improve after a period of overfitting.
[5] 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 [6–9].
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 [10, 11], conservation laws [12]and free energy principle [13].
In particular, phase diagrams and statistical physics have beenidentified as a powerful tool in studying generalization of neural networks [14–17].
However, noneof these works aim to develop an effective theory of representation learning.
We have argued that grokking occurs when models learn a particularly structured representation oftheir inputs.
This connects a low-level understanding of models with their high-level performance.
In a recent work, researchers at Anthropic [18], connect a sudden decrease in loss during trainingwith the emergence of induction heads within their models.
They analogize their work to statisticalphysics , since it bridges a “microscopic”, mechanistic understanding of networks with “macroscopic”facts about overall model performance.
We view our work in the same vein.
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.
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 } 371 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 373 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 378 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]
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 398 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 ≈ [Equation]
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[-
-]
[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]
[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]
[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 = {
[Equation]
[Equation]
[Equation]
[Equation]
[Equation] [Equation] [Equation] (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 472 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]
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]
[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]
h
[Figure]
n
[Figure]
n ) c
[Figure]
m p r
[Figure]
h
[Figure]
n m
[Figure]
m
[Figure]
t
[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 528 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]
s t e p
[Figure]
R Q I >
[Figure]
r c =
[Figure]
(b) 0 .0 .0 .0 .1 .
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)
Acc serves as a lower bound of Acc .
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 to
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).
G Dependence of grokking time on data size
In Section 3.1, we built an effective theory with a quadratic loss function.
ℓ 0 = 12 R T HR , R =[ E 0 , E 1 , · · · , E p − 1 ] , so gradient descent on the loss function is linear, i.e.,
[Equation]
(32) where H has eigenvalues λ i and eigenvectors ¯ v i .
With the initial condition R ( t = 0)
=
a i ¯ v i , wehave R ( t ) =
a i ¯ v i e − λ i t , so the convergence time for the i th eigen-direction is 1 /λ i .
Because the lossfunction is invariant under translation ( E i → E i + c ) and linear scaling ( E i → a E i ) , the first two 571 eigenvalues are zero.
As a result, the time scale towards the linear structure E k = a + k b is 1 /λ 3 .
We call λ 3 the grokking rate.
Note that ℓ 0 involves averaging over parallelograms in the trainingset, it is dependent on training data size, so is λ 3 .
In Figure 17 (a), we plot the dependence of λ 3on training data fraction.
There are many datasets with the same data size, so λ 3 is a probabilisticfunction of data size.
Two insights on grokking can be extracted from this plot: (i) When the data fraction is below somethreshold (around 0.4), λ 3 is zero with high probability, corresponding to no generalization.
This
[Figure]
r a t e
[Figure]
again verifies our critical point in Figure 4.
is implies larger, an important observation from [
To verify our effective theory, we compare the grokking steps obtained from real neural networktraining (defined as steps to RQI > 0 . 95 ), and those predicted by our theory t th ∼ 1 λ 3 η ( η is theembedding learning rate), shown in Figure 17 (b).
The theory agrees qualitatively with neuralnetworks, showing the trend of decreasing grokking steps as increasing data size.
The quantitativedifferences might be explained as the gap between our effective loss and actual loss.
H Effective theory for image classification
In this section, we show our effective theory proposed in Section 3.1 can generalize beyond algorithmic datasets.
In particular, we will apply the effective theory to image classifications.
We find that:(i)
The effective theory naturally gives rise to a novel self-supervised learning method, which canprovably avoid mode collapse without contrastive pairs.
(ii)
The effective theory can shed light on theneural collapse phenomenon [23], in which same-class representations collapse to their class-means.
We first describe how the effective theory applies to image classification.
The basic idea is again that,similar to algorithmic datasets, neural network try to develop a structured representation of the inputsbased on the relational information between samples (class labels in the case of image classification,sum parallelograms in the case of addition, etc.).
The effective theory has two ingredients: (i) sampleswith the same label are encouraged to have similar representations; (ii) the effective loss functionis scale-invariant to avoid all representations collapsing to zero (global collapse).
As a result, aneffective loss for image classification has the form
[Equation]
where x is an image, f ( x ) is its representation, ( x , y ) ∈ P refers to unique pairs x and y that havethe same label.
Scale invariance means the loss function ℓ eff does not change under the linear scalingf ( x ) → a f ( x ) .
Relation to neural collapse It was observed in [23] that image representations in the penultimatelayer of the model have some interesting features: (i) representations of same-class images collapseto their class-means; (ii) class-means of different classes develop into an equiangular tight frame.
Oureffective theory is able to predict the same-class collapse, but does not necessarily put class-meansinto equiangular tight frames.
We conjecture that little explicit repulsion among different classes canhelp class-means develop into an equiangular tight frame, similar to electrons developing into lattice structures on a sphere under repulsive Coulomb forces (the Thomson problem [24]).
We would liketo investigate this modification of the effective theory in the future.
Experiment on MNIST We directly apply the effective loss Eq.
(33) to the MNIST dataset.
Firstly,each image x is randomly encoded to a 2D embedding f ( x ) via the same encoder MLP whose weightsare randomly initialized.
We then train these embeddings by minimizing the effective loss ℓ eff withan Adam optimizer ( 10 − 3 learning rate) for 100 steps.
We show the evolution of these embeddings in
Figure 18.
Images of the same class collapse to their class-means, and different class-means do notcollapse.
This means that our effective theory can give rise to a good representation learning methodwhich only exploits non-contrastive relational information in datasets.
Link to self-supervised learning Note that ℓ itself is vulnerable to global collapse, in the contextof Siamese learning without contrastive pairs.
Various tricks (e.g. decoder with momentum, stopgradient) [8, 25] have been proposed to avoid global collapse.
However, the reasons why thesetricks can avoid global collapse are unclear.
We argue ℓ fails simply because ℓ → a 2 ℓ under scalingf ( x ) → a f ( x ) so gradient descent on ℓ encourage a → 0 .
Based on this picture, our effective theoryprovides another possible fix: make the loss function ℓ scale-invariant (by the normalized loss ℓ eff ), so the gradient flow has no incentive to change representation scales.
In fact, we can prove thatthe gradient flow on ℓ eff preserve Z (variance of representations) so that global collapse is avoidedprovably:
[Equation]
where we use the fact that
[Equation]
I Grokking Experiment on MNIST
[Figure]
- 8.86e2.64e7.85e+
W e i g h t
D e c a y
M e m o r
i z a t
i o n
Grokking
C o m p r e h e n s i o n
C o n f u s i o n (b)
We now demonstrate, for the first time, that grokking (significantly delayed generalization) is a moregeneral phenomenon in machine learning that can occur not only on algorithmic datasets, but also onmainstream benchmark datasets.
In particular, we exhibit grokking on MNIST and demonstrate thatwe can control grokking by varying optimization hyperparameters, just as we did in the main text foralgorithmic datasets in both the full case (transformers) and in the toy model.
To induce grokking on MNIST, we make two nonstandard decisions: (1) we reduce the size of thetraining set from 50k to 1k samples (by taking a random subset) and (2) we increase the scale of theweight initialization distribution (by multiplying the initial weights, sampled with Kaiming uniforminitialization, by a constant > 1 ).
We will provide theoretical motivation for this decision, of changing the initialization distributionto induce grokking, in a follow-up work.
Relevant to this, initialization scale is found to regulate“kernel” vs “rich” learning regimes in networks [26].
{+With these modifications to training set size and initialization scale, we train a depth-3 width-
+}
MLP with ReLU activations with the AdamW optimizer.
We use MSE loss with one-hot targets, rather than cross-entropy.
With this setup, we find that the network quickly fits the train set, and thenmuch later in training validation accuracy improves, as shown in Figure 19a.
This closely follows thestereotypical grokking learning, first observed in algorithmic datasets.
With this setup, we also compute a phase diagram over the model weight decay and the last layerlearning rate.
See Figure 19b.
While in MLPs it is less clear what parts of the network to considerthe “encoder” vs the “decoder”, for our purposes here we consider the last layer to be the “decoder”and vary its learning rate relative to the rest of the network.
The resulting phase diagram has somesimilarity to Figure 6.
We observe a “confusion”phase in the bottom right (high learning rate andhigh weight decay), a “comprehension” phase bordering it, a “grokking” phase as one decreasesweight decay and decoder learning rate, and a “memorization“ phase at low weight decay and lowlearning rate.
Instead of an accuracy threshold of 95%, we use a threshold of 60% here for validationaccuracy for runs to count as comprehension or grokking.
This phase diagram demonstrates that withsufficient regularization, we can again “de-grok” learning.
We also investigate the effect of training set size on time to generalization on MNIST.
We find a resultsimilar to what Power et al.
[1] observed, namely that generalization time increases rapidly once onedrops below a certain amount of training data.
See Figure 20.
[Figure]
J Lottery Ticket Hypothesis Connection
[Figure]
Generalization PCA−−
G e n e r a li z a t i o n
P C A
After generalization − 5 . 0 − 2 .
5 Initialization PCA−−−
I n i t i a li z a t i o n
P C A
At initialization −−  4 Generalization PCA−−−
G e n e r a li z a t i o n
P C A
At initialization
In Figure 21, we show the projection of the learned embeddings after generalization to their firsttwo principal components.
Compared to the projection at initialization, structure clearly emerges inembedding space when the neural network is able to generalize ( > 99% validation accuracy).
What is intriguing is that the projection of the embeddings at initialization to the principal componentsof the embeddings at generalization seems to already contain much of that structure.
In this sense,the structured representation necessary for generalization already existed (partially) at initialization.
The training procedure essentially prunes other unnecessary dimensions and forms the requiredparallelograms for generalization.
This is a nonstandard interpretation of the lottery ticket hypothesiswhere the winning tickets are not weights or subnetworks but instead particular axes or linearcombinations of the weights.
In Figure 22, we show the original training curves (dashed lines).
In solid lines, we recomputeaccuracy with models which use embeddings that are projected onto the ten principal components ofthe embeddings at the end of training (and back).
Clearly, the first few principal components containenough information to reach 99% accuracy.
The first few PCs explain the most variance by definition,however, we note that this is not necessarily the main reason for why they can generalize so well.
Infact, embeddings reconstructed from the PCA at the end of training (solid lines) perform better thancurrent highest variance axes (dotted line).
This behavior is consistent across seeds.
K Derivation of the effective loss
In this section, we will further motivate the use of our effective loss to study the dynamics ofrepresentation learning by deriving it from the gradient flow dynamics on the actual MSE loss inlinear regression.
The loss landscape of a neural network is in general nonlinear, but the linear casemay shed some light on how the effective loss can be derived from actual loss.
For a sample r (whichis the sum of two embeddings E i and E j ), the prediction of the the linear network is D ( r ) =
Ar + b .
The loss function is ( y is its corresponding label):
[Equation]
36) where the first and the second term are prediction error and regularization, respectively.
Both themodel ( A , b ) and the input r are updated via gradient flow, with learning rate η A and η x , respectively:
[Equation]
[Equation]
Inserting ℓ into the above equations, we obtain the gradient flow:
[Equation]
For the d b /dt equation, after ignoring the Ar term and set the initial condition b (0) = 0 , we obtainanalytically b ( t ) =
(1 − e − 2 η A t ) y .
Inserting this into the first and third equations, we have
[Equation]
internal
[Equation]
external force
[Equation]
For the second equation on the evolution of d r /dt , we can artificially decompose the right hand sideinto two terms, based on whether they depend on the label y .
In this way, we call the first term"internal interaction" since it does not depend on y , while the second term "external force".
Notethis distinction seems a bit artificial from a mathematical perspective, but it can be conceptuallyhelpful from a physics perspective.
We will show below the internal interaction term is important forrepresentations to form.
Because we are interested in how two samples interact, we now consider another sample at r ′ , and the evolution becomes
[Equation]
[Equation]
Subtracting d r /dt by d r ′ /dt and setting r ′ = − r , the above equations further simply to
[Equation]
The second equation implies that the pair of samples interact via a quadratic potential U ( r ) =12 r T A T Ar , leading to a linear attractive force f ( r ) ∝ r .
We now consider the adiabatic limit whereη A → 0 .
{+The adiabatic limit Using the standard initialization (e.g., Xavier initialization) of neural networks,we have A T 
+}
{+A 
+}
≈ I .
As a result, the quadratic potential becomes U ( r ) = 12 r T r , which is timeindependent because η A → 0 .
We are now in the position to analyze the addition problem.
For twosamples x (1) =
E i + E j and x (2) =
E m + E n with the same label ( i + j = m + n ), they contributeto an interaction term
[Equation]
[Equation]
Summing over all possible quadruples in the training dataset D , the total energy of the system is
[Equation]
[Equation]
{+where P 
+}
( D ) =
{ ( i, j, m, n ) | i + j = m + n, ( i, j ) ∈ D, ( m, n ) ∈ D } .
To make it scale-invariant, 707 we define the normalized Hamiltonian Eq.
(43) as
[Equation]
which is the effective loss we used in Section 3.1.
[Figure]
[Figure]