Grokking in Linear Estimators -- A Solvable Model that Groks without Understanding
Authors: Noam Itzhak Levi
Alon Beck
Yohai Bar-Sinai
Venue: ICLR 2024 poster (2024)
Keywords: Grokking
Random Matrix Theory
Linear Regression
Representation Learning
Primary area: learning theory
TL;DR: We present and analyze a solvable linear estimation model which exhibits grokking. We find that late generalization increase may not imply a transition from "memorization" to "understanding", but can simply be an artifact of the accuracy measure.

Abstract: Grokking is the intriguing phenomenon where a model learns to generalize long after it has fit the training data. 
    We show both analytically and numerically that grokking can surprisingly occur in linear networks performing linear tasks in a simple teacher-student setup. In this setting, the full training dynamics is derived in terms of the expected training and generalization data covariance matrix. We present exact predictions on how the grokking time depends on input and output dimensionality, train sample size, regularization, and network parameters initialization.
    The key findings are that late generalization increase may not imply a transition from "memorization" to "understanding", but can simply be an artifact of the accuracy measure.
    We provide empirical verification for these propositions, along with preliminary results indicating that some predictions also hold for deeper networks, with non-linear activations.

Published as a conference paper at ICLR 2024
GROKKING IN LINEAR ESTIMATORS - A SOLVABLE MODEL THAT GROKS WITHOUT UNDERSTANDING
Noam Levitt Alon Beck & Yohai Bar Sinai 'Raymond and Beverly Sackler School of Physics and Astronomy Tel-Aviv University
Tel-Aviv 69978, Israel
'écoe Polytechnique Fédérale de Lausanne (EPFL)
Switzerland
(ncamêtau) wateeaau.cc.o ac. alonbk28gmail. com, ybarsinal@gmail com]
ABSTRACT
Grokking the intriguing phenomenon where a model learns to generalize long after it has fit the training data. We show both analytically and numerically that grokking can surprisingly occur linear networks performing linear tasks in a simple teacher-studed setup with Gaussian inputs. In this setting, the full training dynamics is derived in terms the training and generalization data covariance ma- trix. We present exact predictions output dimensionality. train sample size, regularization, and network initialization. We demonstrate that the sharp increase in generalization accuracy may imply a transition from "memorization" to "understanding" ", but can simply be an artifact the accuracy measure. We provide empirical verification for our calculations, along with preliminary results indicating that some predictions also hold for deeper networks, with non-linear activations.
how the grokking time depends input and
INTRODUCTION
Understanding the underlying correlations in complex datasets is the main challenge of statistical learning. Assuming that training and generalization data are drawn from similar distribution, the discrepancy between training and generalization metrics quantifies how well model extracts meaningful features from the training data, and what portion of its reasoning is based on idiosynerasies in the training data. Traditionally, one would expect that once a neural network (NN) training converges to low loss value, the generalization error should either plateau, for good models, or deteriorate for models that overfit.
Surprisingly, Power al. (2022) found that a shallow transformer trained on algorithmic datasets features drastically different dynamics. The network first overfits the training data, achieving low and stable training loss with high generalization error for an extended period, then suddenly and rapidly transitions to perfect generalization phase.
This counter-intuitive phenomenon, dubbed grokking, has recently garnered much attention and many underlying mechanisms have been proposed as possible explanations. These include the difficulty of representation learning (Liu et al., 2022), the scale of parameters at initialization (Liu et al., 2023), spikes in loss ("slingshots") (Thilak al., 2022), random walks among optimal solutions (Millidge, 2022), and the simplicity of the generalising solution (Nanda et al., 2023, Appendix E).
In this paper take different approach, leveraging the simplest possible models that still display grokking- linear estimators. Due to their simplicity, this class of models offers analytically tractable dynamics, allowing derivation of exact predictions for grokking, and clear interpretation that is corroborated empirically. Our main contributions are:
We solve analytically the gradient-fior flow training dynamics in linear teacher-student (T' S E R4 xdout) model performing MSE classification. In this setting, the training and generalization losses Lur, fegns are simply given by ||T -7--Spp S11 where the norm is

---

Published as a conference paper at ICLR 2024
defined with respect to the training/generalization Gram matrices, Exr and Egen respectively, These matrices can be modeled with classical Random Matrix Theory (RMT) techniques.
Grokking in this setting does not imply the simple fact that the generalization loss decays slower than the training loss, because the gradients are set by the latter. The grokking time is mainly determined by a single parameter, the ratio between input dimension and number of training samples À = dia/ Ner-
i any "interesting" generalization behavior, but rather
Standard variations are included in the analysis:
The effect of different weight initializations to generate an artificial rescaling of the training and generalization losses, increasing the effective accuracy value required for saturation and therefore increasing grokking time.
For small dout. Grokking time increases with output dimension due to effectively slower dynamics. This happens up critical dimension after which the measure of accuracy becomes insensitive to the value of the loss, reducing the grokking time,
L2 regularization suppresses grokking in overparameterized networks as expected, while having subtle effect on the grokking time in underparameterizeed settings.
We further show semi-anale aani-anllytically that our results extend to architectures beyond shallow linear networks, including one hidden layer, with both linear and some nonlinear activations.
RELATED WORK
Grokking Many works have attempted to explain the underlying mechanism responsible for grokking, since its discovery by Power et al. (2022). Some works suggest "slingshots" (Thilak et al., 2022) or "oscillations" (Notsawo et al., 2023) underlie grokking. but our explanation applies even without these dynamics. Other works identify ingredients for grokking (Davies et al., 2023; Nanda al., 2023), analyze the trigonometric algorithms networks learn after grokking (Nanda et al., 2023; Chughtai et al., 2023; Merrill et al., 2023), and show similar dynamics in sparse parity tasks (Merrill al., 2023). The addition of regularization has been shown to strongly affect grokking in certain scenarios (Power et al., 2022: Liu et al., 2023). This connection may be attributed to weight decay (WD), for instance, improving generalization (Krogh and Hertz, 1991), though this property is not yet fully understood (Zhang et al., 2018). We incorporate WD in our setup and study its effects on grokking analytically, showing that it can either suppress or enhance grokking, depending on the number of network parameters and number of training samples.
Key related works, most closely tied with our own, are Liu et al. (2022; 2023) and Žunkovič and Ilievski (2022); Gromov (2023). Liu et al. (2022) show perfect generalization on a non-modular addition task when enough data determines the structured representation. Liu et al. (2023) relate grokking to memorization dynamics. Zunkovič and Ilievski (2022); Gromov (2023) analyze solvable models displaying grokking and relate results to latent-space structure formation. Our work employs similar setup but derives grokking dynamics from a random matrix theory perspective relating dataset properties to the empirical covariance matrix.
Linear Estimators in High Dimensions A growing body of work has focused on deriving exact solutions for linear estimators trained on Gaussian data, particularly in the context of random feature models. The dynamics are often described in the gradient flow limit, which we employ in this work. Building on statistical physics methods, Sompolinsky et al. (1988); Advani and Saxe (2017) provided an analytical characterization of the dynamics of learning in linear neural networks under gradient descent, both in shallow and in deep networks Saxe et al. (2014). Their mean-fiet field analysis precisely tracks the evolution of the training and generalization errors, similar to Richards al. (2021); Mignacco et al. (2021): ); Mignacco and Urbani (2022); Paquette et al. (2022). More recently, Bodin and Macris (2022) further studied the dynamies of generalization under gradient descent for the Gaussian covariate model, corroborating the presence of epoch-wire descent structures. In the context of least squares estimation and multiple layers, Loureiro et al. (2022); Goldt et al. (2020) analyzed the gradient flow dynamics and long-time behavior of the training and generalization errors.
The tools from random matrix theory and statistical mechanics employed in these analyses allow precise tracking of the generalization curve and transitions thereof, akin Dobriban and Wager (2015). Our work adopts similar theoretical framing to study the interplay between model capacity, overparameterization, and gradient flow optimization in determining generalization performance.
2

---

Published as a conference paper at ICLR 2024
3 TRAINING DYNAMICS IN A LINEAR -TUUENT SETUP
The majority of our results are derived for a simple student-teacher model Seung et al. (1992), where the inputs are identical independently distributed (iid) normal variables. We draw Nir training samples from a standard Gaussian distribution 2; ~ ~N(0. xd..). and the teacher model generates output labels. The student is trained to mimic the predictions of the teacher, which we take as perfect.
The teacher and student models, which we denote by T and S respectively, share the same architecture. As we show below, Grokking can occur even for the simplest possible network function, which is a linear Perceptron with no biases, or in other words the standard MSE loss. Our analyses are done in the regime of large input dimension and large sample size, i.e., din Nur + oo, where the ratio X = din/No E R kept constant.
simple linear transformation. The loss function
Following the construction presented in Liu al. (2023), we can convert this regression problem into classification task by setting threshold >> 0 and defining a sample to be correctly classified if the prediction error is less than e. The student model is trained with the full batch Gradient Descent (GD) optimizer for steps with a learning rate 77, which may also include weight decay parameter ry. The training loss function is given by
Ner
Lir Nortout
Σ 2m
IIS - TT"
1 -T'(DEEN dout
D=S-T.
(I)
where S,T E Rduxxdou are the student and teacher weight matrices, Exr = 1 En xąaa is the din dmn empirical data covariance, or Gram matrix for the training set, and we define D as the difference between the student and teacher matrices. The elements of T and are drawn at initialization from a normal distribution Sn, T ~ NO, 1/(2din)out) We do not include biases in the student or teacher weight matrices, as they do not affect centrally distributed data.
N.,
Similarly, the generalization loss function is defined as its expectation value over the input distribution, which can be approximated by the empirical average over Ngon randomly sampled points
Lgen EREN
1 dout
|(S -(S T)T--
=|
1 Tr DJ dout
D
(2)
dout
Here Egen is the covariance of the generalization distribution, which the identity. Note that in practice the generalization loss computed by sample average over an independent set, which is not equal to the analytical expectation value. The gradient descent equations at training step / are
Dhr
Der D.
D = (1 (1-
21
Ene
Dr
MY
(D;
(Dz +T),
(3)
dent.
dout
dont
where 7 E + is the weight decay parameter, and RRduxdd is the identity.
It is worthwhile to emphasize the difference between Eq. (1) and Eq. (2), since the distinction between analyses. In training, Eq. (1),
sample average and analytical expectation value is crucial compute the loss over fixed dataset whose covariance, Etr. is non-trivial. The generalization loss defined as the expectation value over the input distribution, which has trivial covariance by assumption, Egen = I. Even if it is computed in practice by averaging over a finite sample with a non-trivial covariance, it is independent of the training dynamics and the sample average will converge to the analytical expectation with the usual VN scaling. This is not true for the training loss, since the training dynamics will guide the network in direction that minimizes the empirical loss with respect to the fixed covariance E,,, This assertion is numerically verified below, as we compare the generalization loss, practically computed by sample averaging, to the analytical result of Eq. (2).
3.1 WARMUP: THE SIMPLEST MODEL.
3.1.1 TRAIN AND GENERALIZATION LOSS
Before analyzing the dynamics of the general linear model, start with simpler setting that captures the most important aspects of the full solution. Concretely, here we set dout = 1, reducing S,T T E Rd from matrices to vectors, and assume weight decay 7 0.
3

---

the two Published as a conference paper at ICLR 2024
Loss
Accuracy
Af grok
10f
106
1.9
10"0
Train
10"
10-9
Test
0.6
10P
Ammly
0.5
10-4
0.1
1.5
10
0.2
10'
Analytic
10"a
100
10P
10
107 107 102
10P
0.0
0.2
0.5
0.5
1.0
Epoclas
Epoclas
Figure 1: Grokking as function Left: Empirical results for training (dashed) and test (solid) losses, for
0.1, 0.9, against analytical solutions (black). For 0.1 test and train curves are indistinguishable Center: Similar comparison for the accuracy functions. Right: Grokking time as a function of A, for different values of the threshold parameter (red, blue), and shown the analytic solution Eq. (11) (dashed black). Here we GD with
Njo 0.01, dn = 100 dout 1 = 10-
Eq. (3) can solved in the gradient flow limit of continuous time, setting 17 = 77tt and + 0, resulting in
((t) = -20
D(I) c
2no Do:
where Do is simply the difference between teacher and student vectors initialization. It follows that the empirical losses, calculated over a dataset admit closed-form expressions as
Lir D Am Eee Do.
Lagen
DJ En Do-
(5)
These expressions for the losses are exact. To proceed, we need to know the Gram matrix of the training dataset, which is the empirical covariance of random sample of Gaussian variables. It known that eigenvectors of Σtr are uniformly distributed on the unit sphere Rd and its eigenvalues, Vi, follow the Marchenko-Pastue (MP) distribution (Marčenko and Pastur, 1967).
1 PMP( = (1 (1 50
v(+ ( X_),
(6)
2T
where 6, is the Dirac mass at R. we define x max( xx 0} for R, and x= (1 vaj².
Since the directions of both D and the eigenvectors of Etr are uniformly distributed, we make the approximation that the projection of on all eigenvectors is the same, which transforms Eq. (5) to the simple form
Lir
din -ee 4rouet
Lyen
din Σ 4rouet
is seen that these sums are the empirical average over the function e en if v follows the MP distribution. This can be well approximated by their respective expectation values,
Cue X, t) Re (X-) [re-Inove
Egen (70, X, s EE-MP(()) [e-1nort]. (8)
The evolution of these loss functions is dictated by the MP distribution, which exhibits distinct behaviors for
For À < the first term in Eq. (6) vanishes, the distribution has no null eigenvalues and so Ler, Leeen both are driven to at f - co, implying that perfect generalization always obtained eventually. On the other hand, for À > Eq. (6) develops several zero eigenvalues, corresponding to flat directions in the training Gram matrix. In this case, while C., is driven to 0, since we -4gort I==0 = O, the generalization loss Cyen does not vanish, as e 4rpat ==u = 1 contributes 1/A À to the loss, preventing perfect generalization. When dout regimes correspond to underparameteriizttoo (A < 1) and overparameterization (A > 1).
and
nonzero constant
In Fig. 1 we show that these analytical predictions are in excellent agreement with numerical experiments, with no fitting parameters, in both regimes.
We also note that the expectation value of Lur in Eq. (8) admits closed form solution, Ltr 4- oF (2; () where FF (a: F (a; (a) is the regularized confluent hyper- geometric function. We could not find closed form expression for Egen, but approximate expressions for the expectation value can be derived for the late time behavior, ef. Appendix B.
4

---

Published as a conference paper at ICLR 2024
3.1.2 TRAIN AND GENERALIZATION ACCURACY
Next, we describe the evolution of the training and generalization accuracy functions. As described above, in the construction of Liu et al. (2023) the accuracy A is defined as the (empirical) fraction of points whose prediction error is smaller than A $ Σ O(e (Dz ()),,))), where Q is the Heaviside step function. We define 2 D R. which is normally distributed with standard deviation D'ED = L, where the covariance of x (that is, Enr for training and I for generalization). Then, in the limit of large sample sizes, the empirical averages converge
2P (1=1 S v Erf V
C 2C
N-000
where A, stand for train and generalization measures. This result implies that the increase in accuracy in the late stages of training can be simply mapped to the decrease of the loss below Writing the accuracy as an explicit function the loss allows an exact calculation of the grokking time, and of whether grokking occurs at all.
3.1.3 GROKKING TIME
In this framework, grokking is simply the phenomenon in which Lie drops below before Leen does. To understand exactly when these events happen, in Appendix B we derive approximate results the long time limit, not > VÀ, showing that
exp -470 (1- V
Egen ~ Lix x xoxxx(11 (1-
Cu (6
(10)
We define grokking time as the time difference between the training and generalization accuracies reaching (v(( s 95%, obtained when each loss satisfies L((") = </4. In terms of the loss functions, show in Appendix B that solving for the difference between *gen tr and expanding the result in the limit of E < 1, one obtains an analytic expression for the grokking time difference
log - vx)
At Atgrok Eggn
(11)
vx)= 2-
21j0
Eq. (II) indicates that the maximal grokking time difference occurs near À ~ where the grokking time diverges quadratically as Atgrok(A + 1) 10- log [T- On the other hand, it vanishes for À 0, which means Nrr >> din and Enr approaches the identity, as expected. These predictions are verified in Fig. (rright).
Effects of Initialization and Label Noise: We briefly comment on the effect of choosing a different initialization for the student weights compared the teacher weights, which discussed in Liu et al. (2023). as well as adding training label noise. In the first setup, rescaling the student weights S - aS leads trivial rescaling of both the training and generalization loss functions L - I+a L. which tantamount to choosing a different threshold parameter - Ita leaving the results unchanged. In the case of training label noise y + y + 8 where 8 ~ N/ ()) the student dynamics don't change, but the training loss function would receive constant contribution, proportional to the noise variance o3. as detailed in Appendix C. It is then possible to tune of and the threshold parameter such that the noise can induce grokking precisely when test accuracy saturates. This effect could be wrongly interpreted as "self-correction" (Liu et al., 2022).
3.2 INTERPRETATION AND INTUITION
We conclude this section by summarizing and interpreting the analytical results for the simple 1-layer linear network with scalar output and MSE loss. In this setting, the loss, which is empirical average over a finite sample, is given by the norm of = S T, as measured by the metric defined by the covariance of the sample, L DTED While the generalization covariance is the identity by
5

---

Published as a conference paper at ICLR 2024
Accuracy
Af grok
107
104
107
Train
Hout
10"1
Nout
Andlytic
Aout
700
0.4
105
10-1
0.5
107
0.8
10"6
10 0.0 107
130
107
307
10a
Epochs
Epochs
dout
Figure 2: Effects of the output dimension dous >> grokking. Left: Empirical results for training (dashed) and gencralization (solid) losses, for dout = 1, 50, 700 (bluc, red, violet) against analytical solutions (black), for = 0.9. Center: Similar comparison for the accuracy functions. Right: The grokking time as a function of dout. for different values of x. Different solid curves are numerical solutions for the expressions given in Section 3.3.1. Training is done using GD with Yoo 0.01, dig 10 10-
construction, the train covariance only approaches the identity the limit No. >> dnn- and otherwise follows the Marchenko-Passe distribution.
The training gradients point to a direction that minimizes the training loss, which is || and in the long time limit, it vanishes exponentially. This must imply that the generalization loss, || which is the norm of the same vector but calculated with respect to a different metric, also vanishes exponentially but somewhat slower. Since in this setting the accuracy is a function of the loss, grokking is identified as the difference between the times that the training and generalization losses fall below the fixed threshold c/4. We note that the fact that the accuracy an explicit function of the loss is a useful peculiarity of this model. In more general settings it is not the case, though it is generally expected that low loss would imply high accuracy.
However, it is noteworthy that nothing particularly interesting is happening at this threshold, and the loss dynamics are oblivious to its existence. other words, grokking this setting, as reported previ- ously by Liu et al. (2023), is an artifact of the definition of accuracy and does not represent a transition from "memorization" to "understanding". or any other qualitative increase in any generalization abilities of the network.
Our analysis can be easily extended to include other effects more complicated scenarios, which we detail below. In all these generalizations the qualitative interpretation remains valid.
3.3
VARIANTS
3.3.
THE BFFECT OF dout
We first extend analysis to the case dout The algebra in this case is similar to what was shown in Section 3.1. We provide the full derivation in Appendix D and report the main results here. The loss evolution follows the same functional form Eq. (8), with the replacement 7oo no/ douts and a correction to Eq. (9):
no (c dout"
I dont dost c)
Couati (No, ),t) -C
.x.t)
=
1
2
26
(12)
dout
r
dout
2
Here, I(a,2)= ( e- the incomplete gamma function, and [(2) So e-4---1 gamma function. These relation hold for Agen/tr and Egen/tr separately.
dt the
We note that the expression for A stems from the fact that ||ı||2 || 211² now follows x² distribution and not a normal distribution. somewhat more complicated.
seen that A still an explicit function of L, albeit
The effects dout > can be read from Eq. (12), and are twofold. Firstly, the accuracy rapidly approaches 1 as the output dimension dout increases, for any value of L and This implies that in the limit of d > oo, both training and generalization accuracies must be close to 100% shortly after initialization and grokking occurs. Secondly, the learning rate fio becomes effectively smaller as

---

Published as a conference paper at ICLR 2024
Loss
Accuracy
2 grok
10
1.--
10
No groleiing
Train
Test
0.8
101
10-1
A- 0.2
== 0.5
10
10
107
0.9
04
10-1
0.2-
102
Analytic
10"
10P
10
109
10
10F
102
10 10
10"
Epocias
Epoclas
Figure Effects of weight decay (2) on grokking. Left: Empirical results for training (dashed) and (blue, red, violet) against analytical solutions (black),
generalization (solid) losses, for y = 10 10 for À = 0.9. Center: Similar comparison for the accuracy functions. Right: The grokking time as a function of y, for different values of X. Different solid curves numerical solutions for the expressions given in Section 3.3.1, while the shaded gray region corresponds training/generalizatioon saturation, without perfect generalization. Training is done using GD with 1j = y/o 0.01, din 10², domn 10
dout grows, implying that the overall time scale of convergence for both training and generalization accuracies increases, leading to a higher grokking time. These two competing effects, along with the monotonicity of the loss functions, give rise to a non-monotonic dependence of the grokking time on dout. which attains a maximum at specific value dmax as can be seen in Fig. 2.
3.3.2 THE EFFECT OF WEIGHT DECAY
We consider first the of nonzero WD in the simpler case of dout 1. Incorporating weight decay amounts to adding a regularization term at each gradient descent timestep, modifying Eq. (3) to
D Dr 2n EE + ? D. nyT.
(13)
where y E R+ is the weight decay parameter. The calculations are straightforward and detailed in Appendix E, the result being that Eq. (8) should be modified to read
[(
+ (e-200( (++ b7)
2
1 - Eu-MP(A) L L=
(14)
2 2
where q == v for the training loss q = 1 for generalization. Compared Eq. (8), it is seen that the main effect of WD is a shift in the effective spectrum v + + $7, as can be expected from the second term in Eq. (13). Since y only affects the gradient but not the accuracy, the expression in Eq. (9) of A as function of C, remains unchanged.
is instructive to analyze Eq. (14) separately for the under and overparameterized regimes. When A < 1, the MP distribution has no null eigenvalues, and the losses begin by decaying exponentially. We can study the grokking behavior by examining the late time limit, i.e. + oc, in which the exponential terms decay, and approximating for small < 1, we obtain the asymptotic expressions
log(1+vA)
Lur
Cgen
At grro
(15)
4(1
A)3
2no (1-VX)
This result means that the generalization loss has higher asymptotic value than the training loss. Thus, there is value of c below which perfect generalization cannot be obtained. For above this threshold, WD has no effect, and below it, the grokking time decreases as given by Eq. (15).
In the overparameterized regime, where À > 1, the MP distribution necessarily contains vanishing eigenvalues, which, as shown Fig. 1, cause the generalization loss to plateau. Introducing weight decay changes this picture somewhat, causing the null eigenvalues to be shifted by factor of 7/2 and ensuring that better generalization performance is reached. Still, the late time behavior is the same Eq. (15), following the same arguments discussed above. We note that in this case, the relevant timescale of the generalization loss is determined by 1/7, leading to suppressing grokking, as observed by Liu et al. (2023), though their regime lies in the gray region of Fig. 3.
The grokking time behaviors for various values of y are clearly demonstrated Fig. In addition, Fig. 7 in the appendix shows the combined effects of X, 7 and dout-
7

---

Published as a conference paper at ICLR 2024
4 GENERALIZATIONS
4.1 2-LAYER NETWORKS
Our analysis can be generalized to multi-layer models. Here, we consider the addition of a single hidden layer, where the teacher network function is f(z) TÇ a x)), where To E R xda, T IR Xdgut, a is an entry-wise activation function and d,, the width of the hidden layer. Similarly, the student network is defined by two matrices (t), S1 (t). The empirical training loss reads
Ltr
of ((8)(()))))) Ti (T a:)
(16)
Ntrdout
In this setup the weights are drawn at initialization from normal distributions So(t = 0), To ~ (00, (2dd)))))) and (t = 0), Ti (O, 1/(2dont
As solvable model, we consider first the case of linear activation, af2) = 2, i.e., a two-layer linear network. In this case, can define T Ti E Rduadda the teacher weights are not updated dynamically. Similar to Eqs. (1) and (2), under the definition D(L)= = S) T,, we show in Appendix F that the gradient flow equations for the system
did in the previous sections, since
2no0
0
D(t) dout
o, 8To D(t))
(17)
Here, h(4) = 1/2 (1)) (())))) (18((0))))) Although Eq. (17) describes a set of coupled equations, we note that the solution for hit)) can be simplified when considering the limit of small noo < 1, as we may ignore the time evolution and consider the trace (or kernel) fixed its initialization value,
which = 1/2 for d, >> dout- In this case, the solutions for L are simple modification of the one given in the previous sections, with the replacement 1/o + no/ (2dout 2). Subsequently, the training/generalizattion performance metries are
22-layer (1o0 X, 1Do
2d 2: x.
110
(18)
and Air/gen is a function of Lir/gens respectively, as given Eq. (12). We note that this setup is generically overparameterized for large da, regardless of dout and for any X. In this sense, all of the results previously derived for < 1 should hold, and grokking occurs as discussed in the previous sections. We experimentally verify that Eq. (18) correctly predicts the performance metrics and their dynamics in Fig. (left column).
4.2 NON-LINEAR ACTIVATIONS
The final extension of our work is to consider the network in Section 4.1, but choose nonlinear activation functions for the hidden layer. the limit large di >>> 1, we expect the network begin to linearize, eventually converging to the Neural Tangent Kernel (NTK) regime (Lee et al., 2019). In this regime, the results in Section 4.1 should hold, up to a redefinition of the kernel which depends on the nonlinearity.
In Fig. (right column), we show that the dynamics 2-layer MLP (1000-200-5) with Lanh activations is well approximated by Eqs. (17) and (18), empirically verifying that our predictions hold beyond the linear regimes, in some cases.
5 DISCUSSION
We have shown that grokking, least in the sense of late generalization, can occur in simple linear teacher-student settings and provided explicit analytical solutions for the training and generalization loss and accuracy dynamics during training. The predictions, which strictly apply in the gradient-fil flow limit and for large sample sizes, were corroborated against numerical experiments and provided an excellent description of the dynamics
8

---

Published as a conference paper at ICLR 2024
Linear activation
tanh activation
10-4
10'
10"5
10-6
10"7
Trim
Test
Accuracy
0.6
du 50
0.1
dg
200
01 Analytio
10
10
102
10
101
Epochs
Epochs
Figure 4: 2-Layer network and nonlinearities. Left: Empirical results for training (dashed) and generalization (solid) losses/accuracies (top/bottom), for two-layer MLP (1000-d;-5) with linear activations and d, 50, 200 (blue, red), against analytical solutions (black). Right: Similar results for an identical network with tanh activations. In both cases, training done using full batch gradient descent with rj = 7/0 0.01, dm 1000, dout = -, 10
main point of this work is that grokking in these cases is an artifact of the accuracy metric and does not represent "understanding" in any meaningful way. Moreover, this scenario does not capture some the observed phenomenology of grokking, which sometimes features non-monotonic loss evolution and special weight structures observed in Gromov (2023); Liu et al. (2022), both of which are not present our setup. Therefore, the proposed mechanism is clearly not a general theory of grokking: the closed-form solutions for the dynamics clearly pertain only to this specific setup, and depend on the spectral properties of the covariance. We thus do not expect that they will generalize "as is" in other contexts.
However, parts of the proposed mechanism may be useful in understanding certain aspects of grokking even in realistic scenarios. The core property of the proposed underlying mechanism of grokking the fact that under GD dynamics, the accuracy is simply function of the loss (cf. Eq. (9) and its generalization Eq. (12)). and importantly a rapidly changing function, with threshold that depends on the data covariance (in our case, this is the condition x Σx < e). In this sense, late generalization an artifact We suspect that this core property. as well as some properties of the analytical predictions, may be observed in more realistic scenarios. In Appendix G we present preliminary results regarding this conjecture, in the context of the classic modular addition task discussed in Gromov (2023). It is shown that the proposed picture is consistent with the numerical experiments.
metric and does not represent any dynamical property or special weight structure.
The application of the proposed mechanism to modeling grokking in realistic settings will be studied in future research. In addition, it would be interesting to study how universal are our results when taking into account finite learning rates, realistic data covariance structures, different choices of optimizers or losses, and other effects.
ACKNOWLEDGEMENTS
We thank Boaz Barak, Nadav Cohen and Andrey Gromov for fruitful discussions. YBS was supported by research grant ISF 1907/22 and Google Gift grant. NL would like to thank the Milner Foundation for the award of a Milner Fellowship. This work was initiated in part at Aspen Center for Physics, which is supported by National Science Foundation grant PHY-2210452.

---

Published as a conference paper at ICLR 2024
REFERENCES
M. S. Advani and A. M. Saxe. High-dimensional dynamics of generalization error in neural networks, 2017.
A. Bodin and N. Macris. Gradient flow the gaussian covariate model: exact solution of learning curves and multiple descent structures, 2022.
B. Chughtai, L. Chan, and N. Nanda. A toy model of universality: Reverse engineering how networks learn group operations, 2023.
X. Davies, L. Langosco, and D. Krueger. Unifying grokking and double descent. arXiv preprint, arXiv:2303.06177 06773, 2023.
E. Dobriban and S. Wager. High-dimensional asymptotics of prediction: Ridge regression and classification, 2015.
Goldt, M. S. Advani, M. Saxe, F. Krzakala, and L. Zdeborova. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher student setup. Journal of Statistical Mechanics: Theory and Experiment, 1020112):12410,, dec 2020. doi: (008//1742-5468/abe URL
https:/ /doi..i..///// www.ooo.oo......o.oo.oooo....... 100.10008888991110100 64 6a$2Fabc616eeee
A. Gromov. Grokking modular arithmetic. arXiv prerint, arXiv:2303.02677 2023.
A. Krogh and J. A. Hertz. A simple weight decay can improve generalization. In NIPS, 1991. URL https: / api. semant iescccolaaa org/Corpus ID: : 10137788.
J. Lee, L. Xiao, S. Schoenholz, Y. Bahri, J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Z. Liu, O. Kitouni, N. S. Nolte, E. Michaud, M. Tegmark, and M. Williams. Towards understand- ing grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651-34663 2022.
Z. Liu, Michaud, and M. Tegmark. Omnigrok: Grokking beyond algorithmic data, 2023.
B. Loureiro, C. Gerbelot, H. Cui, S. Goldt, F. Krzakala, M. Mezard, and L. Zdeborova. Learning curves of generie features maps for realistie datasets with teacher-student model. Journal of Statistical Mechanics: Theory and Experiment, 2022 1 14001, nov 2022. doi: 1081088/ 1742- 5468/ac9825. URL http::tt////eee http://ooi.ooooo 10. 10088221881220
Marčenko and L. Pastur. Distribution of eigenvalues for some sets random matrices. Math- ematics of the USSR-Sbornik, 1(4):457, apr 1967. doi:
URL https:tpp...mm....tmm / dx doi www.oooggo.......
SM1 M0MM1666000000 4ABEH001994.
S. Mei, A. Montanari, and P.-M. Nguyen. A mean field view of the landscape of two-layer neural net- works. Proceedings the National Academy of Sciences, 115(33), July 2018. ISSN 1091-6490. doi: 10. 10.1073/pnas. (1..1773pnaa.180059999 /806579911. URL http:://///. /xx doi. org//0..11100 1073 //77/ppssss.. /paa. ..aas.. 1 806579115.
W. Merrill, N. Tsilivis, and A. Shukla. A tale of two circuits: Grokking competition of sparse and dense subnetworks, 2023.
Mignacco and P. Urbani. The effective noise of stochastic gradient descent. Journal of Statistical Mechanics: Theory and Experiment, 008405) aug 2022. doi: 10. 1484//742--4468//8811 URL https::pp//.......p.m / / org/10/g//// 1
742- -40114722--5445 0.6688FFac814
F. Mignacco, F. Krzakala, P. Urbani, and L. Zdeborova. Dynamical mean-field theory for stochastic
gradient descent in gaussian mixture classification. Journal of Statistical Mechanics: Theory
and Experiment, 2021 (12): 24008, dec 2021. doi: 10. //doi. www.ioorgr1..0 http:/100..88821116610
URL https:
B. Millidge.
Grokking 'grokking',
2022.
URL https: httpp::/wwww.bben....
10

---

Published as a conference paper at ICLR 2024
N. Nanda, L. Chan, T. Liberum, J. Smith, and J. Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023.
P. J. T. Notsawo, H. Zhou, M. Pezeshki, I. Rish, and G. Dumas. Predicting grokking long before it happens: A look into the loss landscape of models which grok, 2023.
C. Paquette, E. Paquette, Adlam, and Pennington. Homogenization of sgd in high-dimensions: Exact dynamics and generalization properties, 2022.
A. Power, Y. Burda, H. Edwards, I. Babuschkin, and V. Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02111 2022.
D. Richards, J. Mourtada, and L. Rosasco. Asymptotics of ridge (less) regression under general source condition, 2021.
A. M. Saxe, L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations, 2014.
H. S. Seung, H. Sompolinsky, and N. Tishby. Statistical mechanics of learning from examples. Physical review A, 45(8):6056, 1992.
H. Sompolinsky, A. Crisanti, and H. J. Sommers. Chaos in random neural networks. Phys. Rev. Lett., 61:259-262, Jul 1988. doi: 10. 10.1100//yyyReeeeee.....55. URL httts://1iikk www.inkk.ppp....m org/ doi/10. 1tt///0.110//hyyaaee www1111/////eeeeeee.
61 .59.
V. Thilak, E. Littwin, S. Zhai, O. Saremi, R. Paiss, and J. Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206.048// 2022.
G. Zhang, C. Wang. B. Xu, and R. B. Grosse. Three mechanisms of weight decay regularization. CoRR, abs/1810.12228 2018. URL http::/ttaaaaeeeeeee /arxivv. wwtg/bb///// /1100.111....11 12281.
B. Žunkovië and E. Ilievski. Grokking phase transitions in learning local rules with gradient descent, 2022.

---

Published as a conference paper at ICLR 2024
A EXPERIMENTAL DETAILS
In all of our experiments, employ a teacher-student model with shared architecture for both teacher and student. The training data consists of fixed number of training samples quoted in the main text for each experiment, drawn from normal distribution (0,, J). All experiments are done on MLPs using MSE loss with default definitions employed by PyTorch. The exact details of each MLP depend on the setup and are quoted in the main text. We train with full batch gradient descent, in all instances. We depart from the default weight initialization of PyTorch, using w ~ NO,,, 1/(2 for cach layer, where d is in the input dimension coming from the previous layer and di is the output dimension of the current layer.
B DERIVATION OF THE GROKKING TIME DIFFERENCE
Here, provide the full derivation for the grokking time difference presented in Eq. (11). Our starting point is the exact solution for the training loss in d case for À < given by
Lr 4no()+1)6 oA (2; 16))
(19)
where Fi (a; 2) oF is the regularized confluent hypergeometric function. We also note the relation
dC -
(20)
dt
which we will relate training and generalization loss functions. Since are interested in the late time behavior, where grokking occurs, we expand the training loss for nol >>>>>> VX, which is given at leading order by
4ma -4 (1-
L1 16VT>3/4(mt))
(21)
Plugging in the result of Eq. (21) into Eq. (20) and integrating over time, we find the expression for the generalization loss at late times is given by
(1-VA) r Anot (1-v2) (--V)))
Vole - 4yu (1-vA)2
ggen
(22)
2 2/ T//o//3/11
A333
where za(ax))= = fee dte dta is the incomplete gamma function. Expanding the result further for late times, we arrive at the result quoted in Eq. (10). In Fig. 5, we show the approximate late-time solutions against the exact solutions. The approximations hold quite well even at somewhat early times, and become increasingly more accurate for later epochs.
With the loss functions hand, we turn the grokking time itself. As described in the main text, define the grokking time as the time difference between the training and generalization accuracies reaching (vf(2 R 95%, obtained when each loss satisfies L(t*) = c/4. Solving this equation for each loss separately, in the late time limit, gives the following expressions for the training and generalization times
2 222/3
33/2
6A
15vA
20
2 W
n
3/3
xVA
(23)
(1- v2)
2/3
8110
25/3
3
TX37I
xVX
7X
W
(24)
gen
8n (1- vX
3/2/3
where W(z) is the Lambert function, which solves the equation WeW = 2, also known as the product-log function. As the argument of both training and generalization times large,
12

---

Published as a conference paper at ICLR 2024
0.001
Than Biastt
10-6
Thai Onder
10-9
10 12
10
100
1000
10
105
106
epochs
Figure 5: Exact training and generalization losses against approximate solutions at late times. In pink/light blue, we show the solutions of Eq. (8). In dashed red is Eq. (21), in dashed blue, we show Eq. (22), while dotted-dashed blue is the solution given in the main text, Eq. (10). Clearly, the asymptotic behavior matches the exact solutions. Here, Yju 0.01, À 0.9, dout 1.
can expand the Lambert function to leading order 2 as ((=) log(z). Taking the difference At grok ggn i and expanding leading order in < 1, we obtain the final expression
(1-
(1
4/3
log
1-VX
3
log
Atgrok = Egee
thr
log
1+
2170
(1-VX)
8n (1 - vA
log
2(2-2VX)
3 3/EV3/2/3
(25)
where the second term goes to zero as é 0, quoted in the main text as Eq. (11).
C DERIVATION FOR TRAINING LABEL NOISE
Here, we consider the addition of training label noise for the simplest setup, in which dout = so S,T are vectors, and in the absence of WD. In this case, the noise addition is a scalar label noise variable 5, ~ NN o.). Training with label noise the linear single-layer model gives the average training loss
No. IISS E
Ctr
T)x; &11 Tr(DDT +03
(26)
Nor
while the generalization loss (w/o noise) is unchanged
Ngan
Eyen
Σ [(S - T)xi11 =T(DD))
(27)
N
gen
Since the label noise contribution does not multiply any of the weights, the gradient the training loss with respect to D is the given in Eq. (3), so the training dynamics in the gradient flow limit remain unchanged
D(t) e Ens Do.
(28)
This implies that the training loss is bounded from below by oi while the generalization loss be driven to O, indicating that even imperfect training in this model leads to perfect generalization.
In Fig. show the loss and accuracy evolutions for different values of o3. Clearly, for high noise variance, the training accuracy saturation a poor metric for performance, while test loss and accuracy show perfect generalization for any noise level.
13

---

Published as a conference paper at ICLR 2024
0.100
Ltzsin.ge get
Agrinage gon
0.5
Thain
0.010
Gen
0.001
0.2
10-1
0.3
0.1
0.1
01001
0.100
10
1000
0.001
0.100
10
1000
epochs
epochs
Figure 6: Training and generalization losses (left) and accuracies (right) for different noise label variance. In shades of pink, show the solutions for Eq. (8) and Eq. (9). In light blue, show the generalization
loss/accuracy evolution which does not change as function of label noise. Here, Y/o = 0.01, À 0.9, dous 1,,6 = 0.01.
D DERIVATION FOR dout >
Here, we provide additional details on the derivation of Eq. (12). The starting point the training and generalization loss functions, given by
1 TT ETT dout
Lir
Egen
Tr E D]
dout "ons
-1112
(29)
dout
where S,T E Xdduss Xdous are the student and teacher weight matrices, E,r 1 xaaa is the empirical data covariance, or Gram matrix for the training set, and we define S - T, the difference between the student and teacher matrices. T and are drawn initialization from normal distributions So, T ~ NO, 1/ 12dd dour)). We do not include biases in the student or teacher weight matrices, as they have no effect on centrally distributed data. The gradient descent equations this instance are simply
Nor
-i=1
D++1 (I (1-
2m dout
Der
Dr.
(30)
where the only difference between the dout = 1 case and the equation above is the rescaled learning rate 1) -> n//dout and the dimensions of D.. Since the MP distribution is identical for each column of Dr, the results sum up and are identical to the dout 1 case for the losses, apart from factor of 1/dout and the learning rate rescaling, leading to Eq. (12).
E Loss CALCULATIONS FOR DYNAMICS INCLUDING WEIGHT DECAY
Here, we provide the derivation for Eq. (14). We begin with the definitions the loss function dout. = 1 case
G = D())
(31)
where D(t) S(t T' the difference between the student and the teacher vectors, Eer N. Σ N, is the training covariance matrix, and 7 2 0 is the weight decay parameter. Using the gradient descent equation in the gradient flow limit, Øt -V, obtain from Eq. (3) that
0D
aD
1 {11) D- m. 2n Eir +
(32)
Øt
Multiplying by the integration factor -o and taking the integral, we arrive at 1 D(1) 3 Eer + 111) T 22( D(0) + " Enr + 2 )1 7
1
(33)
We note that now the limiting value of Dit - o0) is not zero, but rather Doc = v (Ear 11 T. Next, wish caleulate fir D( D(1) and Lgen
14

---

Published as a conference paper at ICLR 2024
0.100
210
0.010
0.010
0.001
0.001
10-4
10-4
29 29
0.0 0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
Figure 7: Grokking time phase diagrams. Left: A contour plot of the grokking time difference as function of y, dont- Shades of red indicate shorter grokking time, while blue tones indicate longer grokking time. White regions indicate no grokking, as generalization accuracy does not converge to 95%. Center and Right: Similar phase diagrams for the grokking time difference as a function of 7, À and dout, A, respectively. The results of all three plots are obtained by numerically finding the grokking time, using the definition A(t") = 0.95 and the analytic formulas quoted in the main text. The fixed parameters for these plots 7ju 0.01, = 10-3
D(t) E gen Dtt), where we emphasize that in both cases D(t) is given by Eq. (33) and depends on Etr. As described in the main text, is a good approximation to set Egen to the identity matrix. For convenience, we will write both cases by Lue/gen = Q(I))Q((I) where Q = Etr for the train and Q I (the identity matrix) for the generalization.
We continue by diagonalizing Err: we write M ET P, where M is a diagonal matrix whose eigenvalues follow the MP distribution. Hence, we obtain
Lur/gon D(t) QD()),
(34)
where Q M, for the train, generalization correspondingly. and D(t) is given by
1 2"1 T 1 5 M (M+271 D(1) e- (o1) D(0) +
-1
M 2 271)
T,
(35)
where D(t) D PTT. We notice the form of: f(M)W where V, W are some vectors, and f(M) is some function the diagonal MP matrix. If V, W random vectors in large dimension, we approximate that
that the expression in Eq. (34) involves terms in
{o V S(MO = { f(u)p(u)du
+ W = W
(36)
where [V] is the norm V, and p(u) is the probability density function of the MP distribution. For example, in our case will get that D² (0) 112 f())((u)du (since D(0) S(0) - T). All that is left now is to calculate the expression in Eq. (34) explicitly, using the approximation of Eq. (36). Doing this, at last arrive into
Lir/gen din / [S()) 4n(ut 127?
C- 27
))
qer /tr/genP(u) (37)
u u+
where 9tr = u and Igen By also setting the student initialization and teacher vector norms to [S(0)),1T] ~ 11///1/22000 1/d (as done the main text), we finally get
2n( 17) 17
Lur/gen = 2 /
400 4o(m+ +
9ur Chr/genP/el p())du.
(38)
+ 1
F DERIVATION FOR THE 2-LAYER NETWORK
Here, we provide supplementary details on the derivation of Eq. (16). We consider the addition of single hidden linear layer, where the teacher network function is f(x) (T(1)) (T(O)TT x, where
15

---

Published as a conference paper at ICLR 2024
T()) Rduxd T())
Rd and d, is the width of the hidden layer. Similarly, the student network is defined by two matrices 5(0), 5()), The empirical training loss over sample set (2; f reads
Nir Ntedout
Lir
(())) - (T"))
(39)
1=1
In this setup the weights are drawn at initialization from normal distributions 5 T())
NO 1/ (2dm)))))) 56" T(2) NO0, dz)) Next, we define T T(0) E Rdm Xdout and derive the gradient flow equations for the system
2no0
Enr
(5())
S T) (5
s(1=
= 2no0 (s (5("))EIH+
(() -)) - (40)
dout
dout
defining D, = s s
T, and noting that D. 56 ((1) 1 (ll)
we arrive at the equations
quoted in the main text
h Eer D,
D. 2100
hn = -ooo D) E.,
(41)
dout
Here, = /2
1|2/2 1 + ////2 where H V VoLa is the Hessian matrix and = { D. Although Eq. (17) describes a set of coupled equations, we note that the solution for he can be simplified when considering
the limit small I7o < 1, as we may ignore the time evolution and consider the trace (or kernel) as fixed to its initialization value, which is ho 1/2 for d, >> dout. In that case the loss solutions are simple modification to the ones given in the previous sections, with the replacement !/o - no/ (2d2uu.) Subsequently, the raaning/generalizzaion performance metrics are
I
(dout dout doutE 22xr//
22-layer layer Ur/gen
|
-1 luyer St/gen
770
x,t
Aur/Ron
(42)
2d
G PRELIMINARY EVALUATION OF PREDICTION CONSISTENCY FOR REAL-WORLD SCENARIOS
Here, we begin to explore to what extent can our predictions for the linear models carry over to real-world settings. As an example, we consider the classical modular arithmetic task, discussed in Power et al. (2022), and later in Gromov (2023).
We borrow the notations and construction from Gromov (2023), using the same architecture and definition of the problem. In this setting, grokking takes place and learned features can be understood analytically. We consider a two-layer MLP network without biases, given by
di
1 d dn
1 (x) = din Σ
(L)) kJ
-P()==((1) (1) (z)), (x)
(2) 9h
(x), k
(43)
where d, is the width of the hidden layer, din is the input dimension, and is activation function, which we take to be the preactivation squared d(h) = 112 At initialization the weights are sampled from the standard normal distribution (1)), w/() ~ NO 1). Eq. (43), we have chosen to follow the mean-field parametrization Mei et al. (2018)).
Given this architecture, we then set up the task of modular arithmetic as classification problem. The network is tasked with learning
f(n. = (n + m) mod
(44)
where m, n, p are integers. To this end, we fix and consider additive over Ep. Each input integer is encoded one-hot vector. The output integer is also encoded as one-hot vector. For the task of
In the limit of infinite width, the meanfield parametrization allows for feature learning.
16

---

Published as a conference paper at ICLR 2024
Seacher-Student MSE classification
Modular arithmetic
Accursey
==
0.5
À 0.4
0.6
À 0.6
0.7
0.8
À 0.95
0.9
Train
Train
10
10-+
16
1
10
10
Train loss
Train loss
10
grok
At
0.04
0.05
0.06
0.08
Figure 8: Grokking dependence on À for modular addition task compared against the Teacher-Student task given in the main text. Top row: In the left panel, we show the accuracy as function training loss Teacher-Student task, while in the right panel. we show the same metric for the modular addition task. Different shades indicate different À value for the Teacher-Studee Blue indicates the smallest number of training samples, leading to longer grokking time, while yellow curves correspond more training samples, therefore shorter grokking time differences. Here, we take p 48 and width da = 400. Bottom panel: Grokking time difference as a function of À 2/(ap) for modular addition. The dashed line shows a (A X*) dependence.
task, or different values for modular arithmetic.
learning bivariate functions over 2, the input dimension is dn 2p, the output dimension is p, the total number of points we split the dataset into train Dirain and test Drest subsets, and furmish this setup with the MSE loss function.
in the dataset is p², while the model in Eq. (43) has 3dp parameters. Finally,
was observed Gromov (2023) that grokking in this setup depends on multiple factors, including the training sample fraction (D.rl/[D] = study the behavior of the grokking time Atgrok. defined in Eq. (11) as a function of À = d (p/(ap²) = 2/ (ap)) by changing the number of training samples, keeping the parameter fixed, and note whether there is any hint of the grokking time being function of the loss alone.
[//D/=Ne/NN] Wee/N Nu./p². In parameterization,
In the top row of Fig. 8, we compare the behavior the accuracies as a function of training loss with different training set sizes for the vanilla linear teacher-student model with dout 1, and for the modular addition task with = 482. We use the natural parameterization for the training set size in both cases, i.e., À in the teacher-student model, and a, the fraction of training samples, for modular addition". We further show on the bottom row of Fig. 8, the grokking time difference as a function of We find that there is some indication that the grokking time difference modular addition depends on À as ~ (A e) 2, which is reminiscent of the linear model. Additionally, comparing the accuracy/taining loss curves, one finds that as grokking begins, the linear and the modular losses and accuracies follow similar trends, indicating that perhaps the linear model predictions can be extended to become universal, under the right reparameterization of the problem.
II order expedite convergence used learning of '7 500. This is in principle "large", but does not affect the standard grokking hchavior observed in Gromov (2023), where the learning rate was set to = 100. 'eeee ranges between 0.5 and 0.9, since below 0.5, the network cannot generalize.
17


Reviews:
[{"review_id": "f1SCaB8VOH", "reviewer": "Reviewer_wJux", "summary": "This paper studies delayed generalization in a linear student-teacher setup with diagonal-covariance Gaussian inputs. The authors also offer some extensions to non-linear 2-layer networks (still in the Gaussian setting). Extensive analytical derivations in the gradient flow regime are shown to match numerical simulations closely.", "questions": "- Delayed generalization in many settings can occur after the generalization loss **increases**, whereas the toy setting presented seems to imply monotonic behavior. Could this mean that the phenomenon studied here is qualitatively different from Grokking in the original setting? There is no fixed and precise definition of Grokking, so this is not a big issue, but I think a distinction might be helpful in understanding feature learning dynamics.\n- The effect of label noise does not seem to be completely captured by the theory. I have seen scenarios where the label noise is “corrected” at Grokking time, still reaching 100% accuracy. Specifically, one can take a *Grokking on modular addition* setup and assign random labels at a certain rate. Generalization accuracy can still reach 100% eventually. Do you have any idea why that’s the case? Is there a way for the toy model to accommodate such \"self-correcting\" behavior?\n\n**Nit:**\nThe figures are nice but can be difficult to parse. In particular, I think it would be easier to understand figures where the theoretical predictions are differentiated for the training and generalization curves (e.g., dashed lines for train, solid for generalization with empirical data shown as points with different markers for train/gen).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is very well-written and easy to follow.\n- Corroborates prior observations on the effect of dataset size on Grokking and distills the delayed generalization behavior down to a very simple toy setup.\n- Theoretically and empirically, the toy setting is well explained and studied thoroughly, from a simple 1-D linear setup to various extensions, including a non-linear 2-layer setting.", "weaknesses": "- My biggest worry about the paper is that it does not deviate from the toy setting proposed. While the results are good, and the match between analytical results and empirical behavior is close, the entire investigation leaves out realistic applications and the original setting in which the phenomenon was first observed. Furthermore, it seems like the theory does not cover some phenomena associated with Grokking, and it might just be too far from the original setting to yield practical results.\nThe submission would be substantially stronger with more connections to realistic settings or insights that can transfer across, e.g., tasks, datasets, models, etc.\n\nOverall, I would tend to recommend acceptance because the thorough investigation and the interesting insights on a simplified toy setting outweigh the limited scope.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4kjID7rnuA", "reviewer": "Reviewer_minU", "summary": "This paper proposes a toy model for grokking, a phenomenon where training accuracy rises much before test accuracy. The proposed model is a linear target function on Gaussian data distribution. A student model attempts to estimate parameters with a matching linear model. The authors use random matrix theory (specifically the eigenvalue distribution of the Wishart ensemble) to derive the dynamics of training and test losses under gradient flow. While both train and test losses decrease smoothly in terms of MSE, when plotting a hard classification accuracy, this toy model exhibits a separation of timescales between training accuracy saturation and test accuracy saturation. The authors illustrate that this is entirely due to the difference between the feature covariance matrix on the training distribution and on the test distribution. These results, taken together provide a possible deflationary explanation for grokking, which can occur merely as an artifact of the choice of metric in very simple models, rather than due to a deeper reason related to learning generalizing features at late training time. The authors show that their results can be extended to weight decay, multiple outputs, training a linear neural network with multiple layers, and training deep nonlinear networks in the lazy/kernel regime.", "questions": "**Questions/Comments**\n\n1. Figure 5 Bottom row: Is the reason that the dynamics for the tanh networks are well approximated by linear two layer dynamics because the activation function is linearizable around $\\phi(h) \\sim h$? What if you used a different activation function, for instance, increasing the “gain” of the tanh nonlinearity like $\\tanh(g * h)$ for $g > 1$? For large enough $g$, would we expect the correspondence to the linearized dynamics to break? \n2. Equation 7 is using approximation signs instead of equality. However, I suspect that if one considered the average case error over random teachers where $T \\sim \\mathcal N(0,I)$ then the equation given is exact. Is there any reason to not focus on averaging over random teachers since the data is isotropic?\n3. The authors claim that deep networks in the kernel regime can be described by the same equations except redefinition of the kernel. I believe that the kernel would not only change the eigenvalue spectrum but would also make the train and test dynamics dependent on the decomposition of the teacher T in the eigenbasis of the kernel. This would therefore make the result much less trivial than what is provided here since the test error is not merely a functional of the spectral density.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "This paper has a nice analytically solvable theory of training dynamics which reproduces grokking-like learning curves. The paper uses a basic result in random matrix theory, namely the limiting spectral density of the Marcheko Pastur law to derive the train and test loss dynamics. Due to the model’s simplicity the authors can derive nice asymptotic expressions for the loss at late time and use these to derive a grokking timescale. They show that their theory is predictive in experiments on reasonably large features and datasets sizes.", "weaknesses": "While this paper provides an interesting and exactly solvable toy model of grokking, it has some major defects which must be addressed before I can support acceptance. \n\nFirst, the paper lacks a proper comparison to literature on gradient flow dynamics in linear models (see below). In addition the citations are often incorrect or correspond to nonexistent papers. We outline these issues below in section titled Related Work Issues.  \n\nNext, the comparison of the grokking observed in this model and grokking observed in real networks is still unclear. Do the authors think that the proposed statistical effect is the phenomenon at play in grokking in real networks? Below I list some observed phenomena associated with grokking that this theory does not quite capture\n\n1.  Many works on grokking report development of specialized weight structures in deep networks temporally coincident with the improved generalization error. I suspect this form of grokking would correspond to deviation from a linear model which would be outside the scope of explanation in this work. If the authors are not claiming this is a theory of grokking as it is observed in its original settings (like modular arithmetic in transformers), in what sense is this an explanation of grokking? \n2. Another example of a difference between observed dynamics and the present study is that weight decay does not induce more extreme grokking which was reported in Liu et al 2023. \n3. Sometimes in grokking experiments with deep networks, the test loss (MSE) even increases before later decreasing to its final value like in Davies et al 2023 Figure 1. It is unclear if this can happen in the linear model where the errors in each non-null eigendirections decrease exponentially.  \n\nOn the other hand, it could be that the primary point of the paper is motivating the need for more careful research on grokking in the future. For example, what artifacts due to choice metric should experimenters be wary of? How should we distinguish “improvement in understanding” vs mere statistical noise? Commenting further on this distinction and comparison to phenomena observed in prior works on grokking would greatly improve the paper. \n\n**Issues with Related Work**\n\n1. The authors miss several important works on gradient flow with linear models. Advani & Saxe 2020 (https://www.sciencedirect.com/science/article/pii/S0893608020303117) derive the training and test errors dynamics for the model considered in this paper. Omitting this citation does not do justice to their relevant contribution. In addition, Saxe et al 2013 (https://arxiv.org/abs/1312.6120) derive learning curves for deep linear networks with fixed dataset. \n2. Several other papers which derive asymptotic performance of linear models trained with Gaussian data in the proportional regime with gradient flow.  \n   (a) Mignacco et al 2022 (https://proceedings.neurips.cc/paper/2020/hash/6c81c83c4bd0b58850495f603ab45a93-Abstract.html), Mignacco et al 2022 (https://iopscience.iop.org/article/10.1088/1742-5468/ac841d/meta), \n   (b) Paquette et al 2022 (https://arxiv.org/abs/2205.07069)\n3. The authors also miss the work of Gromov (https://arxiv.org/abs/2301.02679) on the dynamics of grokking in modular arithmetic without regularization. \n4. Several cited papers do not even seem to exist! These include [2] E. Bodin and N. Macris. Dynamics of generalization in learning with gradient descent for piecewise linear neural networks [4] A. Crisanti and H. Sompolinsky. Dynamics of learning in deep linear neural networks: A mean-field approach   [7] S. Goldt, M. M’ezard, F. Krzakala, and L. Zdeborov’a. Modelling the infinite width limit of neural networks with mean field theory \n\nPlease fix these issues.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cIsb4bdjpE", "reviewer": "Reviewer_Eifc", "summary": "This paper uncovers a surprising fact that linear models can manifest the 'grokking' phenomenon within a teacher-student framework. The author qualitatively predicts the precise dependencies of grokking time on input/output dimensions, sample size, regularization (weight decay), and initialization. Remarkably, these predictions hold true for more intricate models in specific settings. The paper also includes a comprehensive set of empirical validations to support the theoretical analysis.", "questions": "In this linear model, the grokking phenomenon is due to the gap of convergence speed between training and generalization loss. Can it also imply the difference in two-layer networks settings (or in NTK regime), or they have a different mechanism?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is notably well-structured and written. Its strengths can be summarized as follows:\n\n- The paper offers a substantial novelty value, as it explores the grokking phenomenon through a theoretical lens, an area with limited prior analysis. The surprising discovery of grokking in certain linear model scenarios (even without weight decay) adds a unique perspective to the literature. This novel approach contributes significantly to our understanding of grokking from the first principle.\n- The authors provide comprehensive theoretical predictions for their proposed linear model. They not only establish the asymptotic order of grokking time but also consider various related factors. Also, the paper extends its analysis to two-layer networks, both with and without nonlinearity. To further strengthen their claims, the authors conduct an array of extensive experiments that effectively illustrate and validate the theoretical findings.", "weaknesses": "Though the paper is of good quality, I believe the paper would improve in the following aspects:\n- The notations are quite dense even for the simplest model at first glance. The paper would benefit from having some explicit definition or intro to the model and all the notations.\n-  The derivation of this paper is concrete. However, all the calculations are in the asymptotic manner. It would be better to see the results be displayed as some rigorous theorem in some limiting case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fJnmlGIzXb", "reviewer": "Reviewer_VxMM", "summary": "This manuscript studies a linear teacher-student model trained with Gradient flow and MSE loss. The authors solve the dynamical training equation which yields to an exponential decay of the initial error with a spectrum of exponents related to the training data covariance matrix. When the input dimension is larger than the number of data-points, this matrix has zero eigenvalues and hence target modes that cannot be learned. Moreover, using Gaussian iid data, this matrix is Wishart, thereby allowing the authors to estimate where the bulk of non-zero spectral modes resides. Following this they find that the asymptotic dependence of the training loss is roughly exponential with a time scale depending on the lower bulk spectrum of the Wishard matrix. In this asymptotic regime, the test loss behaves as the train loss up to a multiplicative factor larger than 1. Hence the time at which the training loss drops below some small threshold (t_{*,train}) is smaller than the time it takes the test loss to do so (t_{*,test}). This delay between t_{*,test} and t_{*,train} is then viewed as a form of Grokking. The authors proceed with extending their results to several other scenarios. Such as networks with vector outputs, weight decay, and two-layer network in the linear/NTK regime.", "questions": "The work is clearly written. I don't have any questions.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors study a timely topic. \n\nThey lay bare some issues with how Grokking is perceived.  \n\nThe manuscript is clearly written and easy to follow.", "weaknesses": "I feel there are two main issues with the current work.\n\n1. Is this Grokking? \n\nGrokking is a somewhat ill-defined phenomenon. However, like many grey area situations, the fact that it is difficult to draw one sharp reason doesn't mean there cannot be a sharp feeling that something does fit. Let me then rationalize several reasons for why I don't think this is grokking: \nA. It is too simple. If this model Groks, everything Groks.   \nB. There aren't any feature learning effects (see for instance Gromov (2023) or Misra (2022)) or beyond NTK effects. \nC. There is only one time scale here, which is (1-\\sqrt{\\Lambda})^{-2}, governing both test and train loss. Hence, looking at the train and validation loss, a practitioner would not be encouraged to do any early stopping here. In contrast, in various Grokking models, there is an over-fitting regime prior to Grokking which can mislead practitioners.  \n\n2. Technical aspects seem quite close to previous results in the literature. \n\nThe training dynamic the authors find is closely related to the e^{-\\Theta t} \\Delta obtained in the NTK work. In fact, up to a standard shift of the inner and outer indices, their data covariance matrix is Theta here. In particular, both have the same non-zero spectrum. I'm sure there are much earlier works which solve this for a linear network. \n\nSimilar factors relating train and test performances, however at infinite time, were also found in Adlahm and Pennington 2020 and https://arxiv.org/abs/2006.09796. Given the expected exponential decay of the discrepancy in linear models, the fact this ratio holds also in the asymptotic dynamics seems expected. At any rate, It'll be good to compare and disentangle the factors the authors find from those.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)