# Grokking in Linear Estimators – A Solvable Model that Groks without
Understanding

Noam Levi, Alon Beck & Yohai Bar Sinai
Raymond and Beverly Sackler School of Physics and Astronomy
Tel-Aviv University
Tel-Aviv 69978, Israel
{noam@,alonback<EMAIL_ADDRESS>

###### 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 with Gaussian inputs. In this setting,
the full training dynamics is derived in terms of the 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 initialization. We demonstrate that the sharp
increase in generalization accuracy may not imply a transition from
"memorization" to "understanding", but can simply be an artifact of 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.

##  1 Introduction

Understanding the underlying correlations in complex datasets is the main
challenge of statistical learning. Assuming that training and generalization
data are drawn from a similar distribution, the discrepancy between training
and generalization metrics quantifies how well a model extracts meaningful
features from the training data, and what portion of its reasoning is based on
idiosyncrasies in the training data. Traditionally, one would expect that once
a neural network (NN) training converges to a low loss value, the
generalization error should either plateau, for good models, or deteriorate
for models that overfit.

Surprisingly, [18] 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 a 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 [10],
the scale of parameters at initialization [11], spikes in loss ("slingshots")
[21], random walks among optimal solutions [15], and the simplicity of the
generalising solution [16, Appendix E].

In this paper we take a different approach, leveraging the simplest possible
models which still display grokking - linear estimators. Due to their
simplicity, this class of models offers analytically tractable dynamics,
allowing a derivation of exact predictions for grokking, and a clear
interpretation which is corroborated empirically. Our main contributions are:

  * •

We solve analytically the gradient-flow training dynamics in a linear teacher-
student ($T,S\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}$) model
performing MSE classification. In this setting, the training and
generalization losses $\mathcal{L}_{\mathrm{tr}},\mathcal{L}_{\mathrm{gen}}$,
are simply given by $||T-S||^{2}_{\Sigma}$, where the norm is defined with
respect to the training/generalization Gram matrices, $\Sigma_{\mathrm{tr}}$
and $\Sigma_{\mathrm{gen}}$ respectively. These matrices can be modeled with
classical Random Matrix Theory (RMT) techniques.

  * •

Grokking in this setting does not imply any “interesting” generalization
behavior, but rather 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
$\lambda=d_{\mathrm{in}}/N_{\mathrm{tr}}$.

  * •

Standard variations are included in the analysis:

    * –

The effect of different weight initializations is 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 $d_{\mathrm{out}}$, Grokking time increases with output dimension
due to effectively slower dynamics. This happens up to a critical dimension
after which the measure of accuracy becomes insensitive to the value of the
loss, reducing the grokking time.

    * –

$L_{2}$ regularization suppresses grokking in overparameterized networks as
expected, while having a subtle effect on the grokking time in
underparameterized settings.

  * •

We further show semi-analytically that our results extend to architectures
beyond shallow linear networks, including one hidden layer, with both linear
and some nonlinear activations.

##  2 Related work

##### Grokking

Many works have attempted to explain the underlying mechanism responsible for
grokking, since its discovery by [18]. Some works suggest "slingshots" [21] or
"oscillations" [17] underlie grokking, but our explanation applies even
without these dynamics. Other works identify ingredients for grokking [5, 16],
analyze the trigonometric algorithms networks learn after grokking [16, 3,
14], and show similar dynamics in sparse parity tasks [14]. The addition of
regularization has been shown to strongly affect grokking in certain scenarios
[18, 11]. This connection may be attributed to weight decay (WD), for
instance, improving generalization [8], though this property is not yet fully
understood [22]. 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 [10, 11] and [23]. [10]
show perfect generalization on a non-modular addition task when enough data
determines the structured representation. [11] relate grokking to memorization
dynamics. [23] analyze solvable models displaying grokking and relate results
to latent-space structure formation. Our work employs a 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, [4]
provided an analytical characterization of the dynamics of learning in linear
neural networks under gradient descent. Their mean-field analysis precisely
tracks the evolution of the training and generalization errors, similar to
[19]. More recently, [2, 1] further studied the dynamics of generalization
under gradient descent for piecewise linear networks and for the Gaussian
covariate model, corroborating the presence of epoch-wise descent structures.
In the context of least squares estimation and multiple layers, [12, 7]
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 to [6]. Our work adopts a
similar theoretical framing to study the interplay between model capacity,
overparameterization, and gradient flow optimization in determining
generalization performance.

##  3 Training dynamics in a linear teacher-student setup

The majority of our results are derived for a simple student teacher model
[20], where the inputs are identical independently distributed (iid) normal
variables. We draw $N_{\mathrm{tr}}$ training samples from a standard Gaussian
distribution $\mathcal{N}(0,{\bm{I}}_{d_{\mathrm{in}}\times
d_{\mathrm{in}}})$, and the teacher model generates output labels. The student
is trained to mimic the predictions of the teacher, which we take to be
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 – a simple linear transformation. The loss function
is the standard MSE loss. Our analyses are done in the regime of large input
dimension and large sample size, i.e.,
$d_{\mathrm{in}},N_{\mathrm{tr}}\to\infty$, where the ratio $\lambda\equiv
d_{\mathrm{in}}/N_{\mathrm{tr}}\in\mathbb{R}^{+}$ kept constant.

Following the construction presented in [11], we can convert this regression
problem into a classification task by setting a threshold $\epsilon>0$ and
defining a sample to be correctly classified if the prediction error is less
than $\epsilon$. The student model is trained with the full batch Gradient
Descent (GD) optimizer for $t$ steps with a learning rate $\eta$, which may
also include a weight decay parameter $\gamma$. The training loss function is
given by

$\displaystyle\mathcal{L}_{\mathrm{tr}}$
$\displaystyle=\frac{1}{N_{\mathrm{tr}}d_{\mathrm{out}}}\sum_{i=1}^{N_{\mathrm{tr}}}\|(S-T)^{T}x_{i}\|^{2}=\frac{1}{d_{\mathrm{out}}}\textrm{Tr}\left[D^{T}\Sigma_{\mathrm{tr}}D\right],$
$\displaystyle D$ $\displaystyle\equiv S-T\ .$ (1)

where $S,T\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}$ are the
student and teacher weight matrices,
$\Sigma_{\mathrm{tr}}\equiv\frac{1}{N_{\mathrm{tr}}}\sum_{i=1}^{N_{\mathrm{tr}}}x_{i}x^{T}_{i}$
is the $d_{\mathrm{in}}\times d_{\mathrm{in}}$ 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 $S$ are drawn at
initialization from a normal distribution
$S_{0},T\sim\mathcal{N}(0,1/({2d_{\mathrm{in}}}d_{\mathrm{out}}))$. We do not
include biases in the student or teacher weight matrices, as they have no
effect on centrally distributed data.

Similarly, the generalization loss function is defined as its expectation
value over the input distribution, which can be approximated by the empirical
average over $N_{\mathrm{gen}}$ randomly sampled points

$\displaystyle\mathcal{L}_{\mathrm{gen}}$
$\displaystyle=\mathbb{E}_{x\sim\mathcal{N}}\left[\frac{1}{d_{\mathrm{out}}}\|(S-T)^{T}x\|^{2}\right]=\frac{1}{d_{\mathrm{out}}}\textrm{Tr}\left[D^{T}\Sigma_{\mathrm{gen}}D\right]=\frac{1}{d_{\mathrm{out}}}\|D\|^{2}\
.$ (2)

Here $\Sigma_{\mathrm{gen}}$ is the covariance of the generalization
distribution, which is the identity. Note that in practice the generalization
loss is computed by a sample average over an independent set, which is not
equal to the analytical expectation value. The gradient descent equations at
training step $t$ are

$\displaystyle\nabla_{D}\mathcal{L}_{\mathrm{tr}}$
$\displaystyle=\frac{2}{d_{\mathrm{out}}}\Sigma_{\mathrm{tr}}D\ ,$
$\displaystyle{D}_{t+1}=\left({\bm{I}}-\frac{2\eta}{d_{\mathrm{out}}}\Sigma_{\mathrm{tr}}\right)D_{t}-\frac{\eta\gamma}{d_{\mathrm{out}}}\left(D_{t}\
+T\right),$ (3)

where $\gamma\in\mathbb{R}^{+}$ is the weight decay parameter, and
${\bm{I}}\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{in}}}$ is the
identity.

It is worthwhile to emphasize the difference between Eq. 1 and Eq. 2, since
the distinction between sample average and analytical expectation value is
crucial to our analyses. In training, Eq. 1, we compute the loss over a fixed
dataset whose covariance, $\Sigma_{\mathrm{tr}}$, is non trivial. The
generalization loss is defined as the expectation value over the input
distribution, which has a trivial covariance by assumption,
$\Sigma_{\mathrm{gen}}={\bm{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 $\sqrt{N}$ scaling. This is not true
for the training loss, since the training dynamics will guide the network in a
direction that minimizes the empirical loss with respect to the fixed
covariance $\Sigma_{\mathrm{tr}}$. 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, we start with a
simpler setting which captures the most important aspects of the full
solution. Concretely, here we set $d_{\mathrm{out}}=1$, reducing
$S,T\in\mathbb{R}^{d_{\mathrm{in}}}$ from matrices to vectors, and assume no
weight decay $\gamma=0$. Eq. 3 can be solved in the gradient flow limit of
continuous time, setting $\eta=\eta_{0}dt$ and $dt\to 0$, resulting in

$\displaystyle\dot{D}(t)=-2\eta_{0}\Sigma_{\mathrm{tr}}D(t)\quad\to\quad
D(t)=e^{-2\eta_{0}\Sigma_{\mathrm{tr}}t}D_{0},$ (4)

where $D_{0}$ is simply the difference between teacher and student vectors at
initialization. It follows that the empirical losses, calculated over a
dataset functions admit closed form expressions as

$\displaystyle\mathcal{L}_{\mathrm{tr}}=D_{0}^{T}e^{-4\eta_{0}\Sigma_{\mathrm{tr}}t}\Sigma_{\mathrm{tr}}D_{0},\qquad\mathcal{L}_{\mathrm{gen}}=D_{0}^{T}e^{-4\eta_{0}\Sigma_{\mathrm{tr}}t}D_{0}.$
(5)

Figure 1:  Grokking as a function of $\lambda$. Left: Empirical results for
training (dashed) and generalization (solid) losses, for $\lambda=0.1,0.9,1.5$
(red, blue, violet) against analytical solutions (black). Center: Similar
comparison for the accuracy functions. Right: The grokking time as a function
of $\lambda$, for different values of the threshold parameter $\epsilon$.
Different solid curves are numerical solutions for the expressions given in
Section 3.1, shown against the analytic solution in Eq. 12 (dashed black). In
all three panels, diamonds/stars indicate training/generalization accuracy
convergence to 95%. Training is done using GD with
$\eta=\eta_{0}=0.01,d_{\mathrm{in}}=\\!10^{3},d_{\mathrm{out}}=1,\epsilon=\\!10^{-3}$.

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 a
random sample of Gaussian variables. It is known that eigenvectors of
$\Sigma_{\mathrm{tr}}$ are uniformly distributed on the unit sphere in
$\mathbb{R}^{d_{\mathrm{out}}}$ and its eigenvalues, $\nu_{i}$, follow the
Marchenko-Pastur (MP) distribution [13],

$\displaystyle
p_{\mathrm{MP}}(\nu)d\nu=\left(1-\frac{1}{\lambda}\right)^{+}\delta_{0}+\frac{\sqrt{(\lambda_{+}-\nu)(\nu-\lambda_{-})}}{2\pi\lambda\nu}I_{\nu\in[\lambda_{-},\lambda_{+}]}d\nu,$
(6)

where $\delta_{\nu}$ is the Dirac mass at $\nu\in\mathbb{R}$, we define
$x^{+}=\max\\{x,0\\}$ for $x\in\mathbb{R}$, and
$\lambda_{\pm}=(1\pm\sqrt{\lambda})^{2}$.

Since the directions of both $D$ and the eigenvectors of
$\Sigma_{\mathrm{tr}}$ are uniformly distributed, we make the approximation
that the projection of $D$ on all eigenvectors is the same, which transforms
Eq. 5 to the simple form

$\displaystyle\mathcal{L}_{\mathrm{tr}}$
$\displaystyle\approx\frac{1}{d_{\mathrm{in}}}\sum_{i}e^{-4\eta_{0}\nu_{i}t}\nu_{i}\
,$ $\displaystyle\mathcal{L}_{\mathrm{gen}}$
$\displaystyle\approx\frac{1}{d_{\mathrm{in}}}\sum_{i}e^{-4\eta_{0}\nu_{i}t}\
.$ (7)

It is seen that these sums are the empirical average over the function
$e^{-4\eta_{0}\nu t}\nu$, if $\nu$ follows the MP distribution. This can be
well approximated by their respective expectation values,

$\displaystyle\mathcal{L}_{\mathrm{tr}}(\eta_{0},\lambda,t)$
$\displaystyle\approx\mathbb{E}_{\nu\sim\mathrm{MP(\lambda)}}\left[\nu
e^{-4\eta_{0}\nu t}\right],$
$\displaystyle\mathcal{L}_{\mathrm{gen}}(\eta_{0},\lambda,t)$
$\displaystyle\approx\mathbb{E}_{\nu\sim\mathrm{MP(\lambda)}}\left[e^{-4\eta_{0}\nu
t}\right].$ (8)

The evolution of these loss functions are dictated by the MP distribution,
which exhibits distinct behaviors for $\lambda\lessgtr 1$. For $\lambda<1$,
the first term in Eq. 6 vanishes, the distribution has no null eigenvalues and
so $\mathcal{L}_{\mathrm{tr}},\mathcal{L}_{\mathrm{gen}}$ both are driven to
$0$ at $t\to\infty$, implying that perfect generalization is always obtained
eventually. On the other hand, for $\lambda>1$, Eq. 6 develops a number of
zero eigenvalues, corresponding to flat directions in the training Gram
matrix. In this case, while $\mathcal{L}_{\mathrm{tr}}$ is driven to 0, since
$\nu e^{-4\eta_{0}\nu t}|_{\nu=0}=0$, the generalization loss
$\mathcal{L}_{\mathrm{gen}}$ does not vanish, as $e^{-4\eta_{0}\nu
t}|_{\nu=0}=1$ contributes a nonzero constant $1-\frac{1}{\lambda}$ to the
loss, preventing perfect generalization. When $d_{\mathrm{out}}=1$, the two
regimes correspond directly to underparameterization ($\lambda<1$) and
overparameterization ($\lambda>1$).

In Fig. 1 we show that these analytical prediction are in excellent agreement
with numerical experiments, with no fitting parameters, in both regimes.

We also note that the expectation value of $\mathcal{L}_{\mathrm{tr}}$ in Eq.
8 admits a closed form solution,

$\mathcal{L}_{\mathrm{tr}}=e^{-4\eta_{0}(\lambda+1)t}\,_{0}\tilde{F}_{1}\left(2;16\eta_{0}^{2}t^{2}\lambda\right),$
(9)

where ${}_{0}\tilde{F}_{1}\left(a;z\right)={{}_{0}{F}_{1}}(a;z)\Gamma(a)$ is
the regularized confluent hypergeometric function. We could not find a closed
form expression for $\mathcal{L}_{\mathrm{gen}}$, but approximate expressions
for the expectation value can be derived for the late time behavior, cf.
Appendix B.

####  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 [11] the accuracy
$\mathcal{A}$ is defined as the (empirical) fraction of points whose
prediction error is smaller than $\epsilon$,
$\mathcal{A}=\frac{1}{N}\sum_{i=1}^{N}\Theta(\epsilon-(D^{T}(t)x_{i})^{2})$,
where $\Theta$ is the Heaviside step function. We define
$z=D^{T}x\in\mathbb{R}$, which is normally distributed with standard deviation
$D^{T}\Sigma D=\mathcal{L}$, where $\Sigma$ is the covariance of $x$ (that is,
$\Sigma_{\mathrm{tr}}$ for training and ${\bm{I}}$ for generalization). Then,
in the limit of large sample sizes the empirical averages converge to

$\displaystyle\mathcal{A}_{\mathrm{tr}/\mathrm{gen}}\to
2\Pr\left(|z|\leq\sqrt{\epsilon}\right)=\mathrm{Erf}\left(\sqrt{\frac{\epsilon}{2\mathcal{L}_{\mathrm{tr}/\mathrm{gen}}}}\right)\
.$ (10)

The implication of this result is that the increase in accuracy in late stages
of training can be simply mapped to the decrease of the loss below $\epsilon$.
Writing the accuracy as an explicit function of 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
$\mathcal{L}_{\mathrm{tr}}$ drops below $\epsilon$ before
$\mathcal{L}_{\mathrm{gen}}$ does. To understand exactly when these events
happen, in Appendix B we derive approximate results in the long time limit,
$\eta_{0}t\gg\sqrt{\lambda}$, showing that

$\displaystyle\mathcal{L}_{\mathrm{tr}}\simeq\frac{e^{-4\eta_{0}\left(1-\sqrt{\lambda}\right)^{2}t}}{16\sqrt{\pi}\lambda^{3/4}(\eta_{0}t)^{3/2}},\quad\mathcal{L}_{\mathrm{gen}}\simeq\mathcal{L}_{\mathrm{tr}}\times\left(1-\sqrt{\lambda}\right)^{-2}.$
(11)

We define grokking time as the time difference between the training and
generalization accuracies reaching $\operatorname{Erf}(\sqrt{2})\approx 95\%$,
obtained when each loss satisfies $\mathcal{L}(t^{*})=\epsilon/4$. In terms of
the loss functions, we show in Appendix B that solving for the difference
between $t^{*}_{\mathrm{gen}}-t^{*}_{\mathrm{tr}}$, and expanding the result
in the limit of $\epsilon\ll 1$, one obtains an analytic expression for the
grokking time difference

$\displaystyle\Delta
t_{\mathrm{grok}}=t^{*}_{\mathrm{gen}}-t^{*}_{\mathrm{tr}}\simeq\tfrac{\log\left(\frac{1}{1-\sqrt{\lambda}}\right)}{2\eta_{0}\left(1-\sqrt{\lambda}\right)^{2}}.$
(12)

Eq. 12 indicates that the maximal grokking time difference occurs near
$\lambda\simeq 1$, where the grokking time diverges quadratically as $\Delta
t_{\mathrm{grok}}(\lambda\to
1)\sim\frac{1}{\eta_{0}(\lambda-1)^{2}}\log\left(\frac{4}{(1-\lambda)^{2}}\right)$.
On the other hand, it vanishes for $\lambda\simeq 0$, which means
$N_{\mathrm{tr}}\gg d_{\mathrm{in}}$ and $\Sigma_{\mathrm{tr}}$ approaches the
identity, as expected. These predictions are verified in Fig. 1(right).

##### Effects of Initialization and Label Noise:

We briefly comment on the effect of choosing a different initialization for
the student weights compared to the teacher weights, which is discussed in
[11], as well as adding training label noise. In the first setup, rescaling
the student weights $S\to\alpha S$ leads to a trivial rescaling of both the
training and generalization loss functions as
$\mathcal{L}\to\frac{1+\alpha^{2}}{2}\mathcal{L}$, which is tantamount to
choosing a different threshold parameter
$\epsilon\to\frac{2\epsilon}{1+\alpha^{2}}$, leaving the results unchanged. In
the case of training label noise $y\to y+\delta$ , where
$\delta\sim\mathcal{N}(0,\sigma_{\delta}^{2})$, the student dynamics don’t
change, but the generalization loss function would receive a constant
contribution, proportional to the noise variance $\sigma^{2}_{\delta}$. This
contribution will simply imply that for small $\epsilon$, grokking to perfect
generalization cannot occur, but rather just to some finite accuracy.

####  3.1.4 Interpretation and intuition

We conclude this section by summarizing and interpreting the analytical
results for the simple 1-layer linear network with a scalar output and MSE
loss. In this setting, the loss, which is an empirical average over a finite
sample, is given by the norm of $D=S-T$, as measured by the metric defined by
the covariance of the sample, $\mathcal{L}=D^{T}\Sigma D$. While the
generalization covariance is the identity by construction, the train
covariance only approaches the identity in the limit $N_{\mathrm{tr}}\gg
d_{\mathrm{in}}$, and otherwise follows the Marchenko-Pastur distribution.

The training gradients point to a direction that minimizes the training loss,
which is $\|D\|_{\Sigma_{\mathrm{tr}}}$, and in the long time limit it
vanishes exponentially. This must imply that the generalization loss,
$\|D\|_{\bm{I}}$, which is the norm of the same vector but calculated wirth
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 $\epsilon/4$. We note
that the fact that the accuracy is 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. In
other words, grokking in this setting, as reported previously by [11], 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 in more
complicated scenarios, which we detail below. In all these generalizations the
qualitative interpretation remains valid.

###  3.2 The effect of $d_{\mathrm{out}}$

We first extend our analysis to the case $d_{\mathrm{out}}>1$. The algebra in
this case is similar to what was shown in Section 3.1. We provide the full
derivation in Appendix C and report the main results here.

The loss evolution follows the same functional form as Eq. 8, with the
replacement $\eta_{0}\to\eta_{0}/d_{\mathrm{out}}$. In addition, when
$d_{\mathrm{out}}>1$ the mapping between $\mathcal{L}$ and $\mathcal{A}$, Eq.
10, should be corrected since $\|z\|^{2}=\|D^{T}x\|^{2}$ now follows a
$\chi^{2}$ distribution and not a normal distribution, resulting in

$\displaystyle\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}^{d_{\mathrm{out}}\geq 1}$
$\displaystyle=\frac{1}{d_{\mathrm{out}}}\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}^{d_{\mathrm{out}}=1}\left(\frac{\eta_{0}}{d_{\mathrm{out}}},\lambda,t\right)\
,$ $\displaystyle\mathcal{A}_{\mathrm{\mathrm{tr}/gen}}$
$\displaystyle=1-\frac{\Gamma\left(\frac{d_{\mathrm{out}}}{2},\frac{d_{\mathrm{out}}\epsilon}{2\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}}\right)}{\Gamma\left(\frac{d_{\mathrm{out}}}{2}\right)},$
(13)

where $\Gamma(a,z)=\int_{z}^{\infty}dte^{-t}t^{a-1}$ is the incomplete gamma
function, and $\Gamma(z)=\int_{0}^{\infty}dte^{-t}t^{z-1}$ is the gamma
function. It is seen that $\mathcal{A}$ is still an explicit function of
$\mathcal{L}$, albeit somewhat more complicated.

The effects of $d_{\mathrm{out}}>1$ can be read from Eq. 13, and are twofold.
Firstly, the accuracy rapidly approaches $1$ as the output dimension
$d_{\mathrm{out}}$ increases, for any value of $\mathcal{L}$ and $\epsilon$.
This implies that in the limit of $d_{\mathrm{out}}\to\infty$, both training
and generalization accuracies must be close to 100% shortly after
initialization and no grokking occurs. Secondly, the learning rate $\eta_{0}$
becomes effectively smaller as $d_{\mathrm{out}}$ 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 $d_{\mathrm{out}}$, which
attains a maximum at a specific value $d_{\mathrm{out}}^{\mathrm{max}}$, as
can be seen in Fig. 2.

Figure 2:  Effects of the output dimension $d_{\mathrm{out}}>1$ on grokking.
Left: Empirical results for training (dashed) and generalization (solid)
losses, for $d_{\mathrm{out}}=1,50,700$ (blue, red, violet) against analytical
solutions (black), for $\lambda=0.9$. Center: Similar comparison for the
accuracy functions. Right: The grokking time as a function of
$d_{\mathrm{out}}$, for different values of $\lambda$. Different solid curves
are numerical solutions for the expressions given in Section 3.2. In all three
panels, diamonds/stars indicate training/generalization accuracy convergence
to 95%, shown for $d_{\mathrm{out}}^{\mathrm{max}}\simeq 50$, where the
grokking time is maximal. Training is done using GD with
$\eta=\eta_{0}=0.01,d_{\mathrm{in}}=\\!10^{3},\epsilon=\\!10^{-3}$.

###  3.3 The effect of weight decay

We consider first the case of nonzero WD in the simpler case of
$d_{\mathrm{out}}=1$. Incorporating weight decay amounts to adding a
regularization term at each gradient descent timestep, modifying Eq. 3 to

$\displaystyle D_{t+1}=D_{t}-2\eta\left(\Sigma_{\mathrm{tr}}+\frac{1}{2}\gamma
I\right)D_{t}-\eta\gamma T,$ (14)

where $\gamma\in\mathbb{R}^{+}$ is the weight decay parameter. Comparing to
Eq. 3, it is seen that this basically amounts to shifting the eigenvalues of
$\Sigma$ by $\gamma$. The calculations are straightforward and detailed in
Appendix D, the result being that Eq. 8 is modified to read

$\displaystyle\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}=\frac{1}{2}\mathbb{E}_{\nu\sim\mathrm{MP}(\lambda)}\left[\left(e^{-4\eta_{0}\left(\nu+\frac{1}{2}\gamma\right)t}+\left(\frac{e^{-2\eta_{0}\left(\nu+\frac{1}{2}\gamma\right)t}\nu+\frac{1}{2}\gamma}{\nu+\frac{1}{2}\gamma}\right)^{2}\right)q_{\mathrm{tr}/\mathrm{gen}}\right],$
(15)

where $q_{\mathrm{tr}}=\nu$ and $q_{\mathrm{gen}}=1$. Since $\gamma$ only
affects the gradient but not the accuracy, the expression in Eq. 10 of
$\mathcal{A}$ as a function of $\mathcal{L}$, remains unchanged.

Figure 3:  Effects of weight decay ($\gamma$) on grokking. Left: Empirical
results for training (dashed) and generalization (solid) losses, for
$\gamma=10^{-5},10^{-3},10^{-2}$ (blue, red, violet) against analytical
solutions (black), for $\lambda=0.9$. Center: Similar comparison for the
accuracy functions. Right: The grokking time as a function of $\gamma$, for
different values of $\lambda$. Different solid curves are numerical solutions
for the expressions given in Section 3.2, while the shaded gray region
corresponds to training/generalization saturation, without perfect
generalization. In all three panels, diamonds/stars indicate the point where
accuracy reaches 95%. Training is done using GD with
$\eta=\eta_{0}=0.01,d_{\mathrm{in}}=\\!10^{3},d_{\mathrm{out}}=1,\epsilon=\\!10^{-3}$
Figure 4:  Grokking time phase diagrams. Left: A contour plot of the grokking
time difference as a function of $\gamma,d_{\mathrm{out}}$. 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 $\gamma,\lambda$ and
$d_{\mathrm{out}},\lambda$, respectively. The results of all three plots are
obtained by numerically finding the grokking time, using the definition
$\mathcal{A}(t^{*})=0.95$ and the analytic formulas quoted in the main text.
The fixed parameters for these plots are $\eta_{0}=0.01,\epsilon=\\!10^{-3}$.

It is instructive to analyze Eq. 15 separately for the under and
overparameterized regimes. When $\lambda<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. $t\to\infty$, in
which the exponential terms decay, and approximating for small $\gamma\ll 1$,
we obtain the asymptotic expressions

$\displaystyle\mathcal{L}_{\mathrm{tr}}\simeq\frac{\gamma^{2}}{4(1-\lambda)},\qquad\mathcal{L}_{\mathrm{gen}}\simeq\frac{\gamma^{2}}{4(1-\lambda)^{3}},\qquad\Delta
t_{\mathrm{grok}}\simeq\tfrac{\log\left(1+\sqrt{\lambda}\right)}{2\eta_{0}\left(1-\sqrt{\lambda}\right)^{2}}.$
(16)

This result means that the generalization loss has a higher asymptotic value
than the training loss. Thus, there is a value of $\epsilon$ below which
perfect generalization cannot be obtained. For $\epsilon$ above this threshold
WD has no effect, and below it the grokking time decreases as given by Eq. 16.

In the overparameterized regime, where $\lambda>1$, the MP distribution
necessarily contains vanishing eigenvalues, which, as shown in Fig. 1, cause
the generalization loss to plateau. Introducing weight decay changes this
picture somewhat, causing the null eigenvalues to be shifted by a factor of
$\gamma/2$ and ensuring that better generalization performance is reached.
Still, the late time behavior is the same as Eq. 16, following the same
arguments as discussed above. We note that in this case, the relevant
timescale of the generalization loss is determined by $1/\gamma$, leading to
suppressing grokking, as noted by [11].

The grokking time behaviors for various values of $\gamma$ are clearly
demonstrated in Fig. 3.

In Fig. 4, we summarize our results by combining all the separate effects,
showing two dimensional slices of the grokking phase diagram, which depends on
$\lambda,d_{\mathrm{out}}$ and $\gamma$, mirroring each separate effect.

##  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(x)=(T^{(1)})^{T}\sigma((T^{(0)})^{T}x)$, where
$T^{(0)}\in\mathbb{R}^{d_{\mathrm{in}}\times d_{h}}$,
$T^{(1)}\in\mathbb{R}^{d_{h}\times d_{\mathrm{out}}}$, $\sigma$ is an entry-
wise activation function and $d_{h}$ is the width of the hidden layer.
Similarly, the student network is defined by two matrices $S^{(0)},S^{(1)}$.
The empirical training loss reads

$\displaystyle\mathcal{L}_{\mathrm{tr}}$
$\displaystyle=\frac{1}{N_{\mathrm{tr}}d_{\mathrm{out}}}\sum_{i=1}^{N_{\mathrm{tr}}}\left((S^{(1)})^{T}\sigma((S^{(0)})^{T}x_{i})-(T^{(1)})^{T}\sigma((T^{(0)})^{T}x_{i})\right)^{2}\
.$ (17)

In this setup, the weights are drawn at initialization from normal
distributions
$S^{(0)}_{0},T^{(0)}\sim\mathcal{N}(0,1/({2d_{\mathrm{in}}}d_{h}))$ and
$S_{0}^{(1)},T^{(1)}\sim\mathcal{N}(0,1/({2d_{\mathrm{out}}}d_{h}))$.

As a solvable model, we consider first the case of linear activation,
$\sigma(z)=z$, i.e., a two layer linear network. In this case we can define
$T=T^{(0)}T^{(1)}\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}$ as we
did in the previous sections, since the teacher weights are not updated
dynamically. Similar to Eqs. 1 and 2, under the definition
$D_{t}=S_{t}^{(0)}S_{t}^{(1)}-T$, we show in Appendix E that the gradient flow
equations for the system are

$\displaystyle\dot{D}_{t}=-2\eta_{0}\frac{h}{d_{\mathrm{out}}^{2}}\Sigma_{\mathrm{tr}}D,\qquad\dot{h}_{t}=-8\eta_{0}(T+D)^{T}\Sigma_{\mathrm{tr}}D.$
(18)

Here, $h=Tr[H]/2=\|S^{(0)}\|^{2}/2+\|S^{(1)}\|^{2}/2$, where
$H=\nabla_{\theta}^{T}\nabla_{\theta}\mathcal{L}_{\mathrm{tr}}$ is the Hessian
matrix and $\theta\equiv\\{S^{(0)},S^{(1)}\\}$. Although Eq. 18 describes a
set of coupled equations, we note that the solution for $h_{t}$ can be
simplified when considering the limit of small $\eta_{0}\ll 1$, as we may
ignore the time evolution and consider the trace (or kernel) as fixed to its
initialization value, which is $h_{0}\simeq 1/2$ for $d_{h}\gg
d_{\mathrm{out}}$. In that case the loss solutions are a simple modification
to the ones given in the previous sections, with the replacement
$\eta_{0}\to\eta_{0}/(2d_{\mathrm{out}}^{2})$. Subsequently, the
training/generalization performance metrics are

$\displaystyle\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}^{\mathrm{2-layer}}$
$\displaystyle=\|D_{0}\|^{2}\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}^{\mathrm{1-layer}}\left(\frac{\eta_{0}}{2d_{\mathrm{out}}^{2}},\lambda,t\right),\quad\mathcal{A}_{\mathrm{\mathrm{tr}/gen}}=1-\frac{\Gamma\left(\frac{d_{\mathrm{out}}}{2},\frac{d_{\mathrm{out}}\epsilon}{2\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}}\right)}{\Gamma\left(\frac{d_{\mathrm{out}}}{2}\right)}.$
(19)

We note that this setup can be generically classified within the
overparameterized regime, provided that $d_{h}\gg 1$, regardless of
$d_{\mathrm{out}}$ and for any $\lambda$. In this sense, all of the results
previously derived for $\lambda<1$ hold, and grokking occurs as discussed in
the previous sections. We experimentally verify that Eq. 19 correctly predicts
the performance metrics and their dynamics in Fig. 5 (top row).

Figure 5:  2-Layer network and nonlinearities. Top row: Empirical results for
training (dashed) and generalization (solid) losses/accuracies (left/right),
for a two layer MLP (1000-$d_{h}$-5) with linear activations and
$d_{h}=50,200$ (blue, red), against analytical solutions (black). Bottom row:
Similar results, for a two layer MLP (1000-$d_{h}$-5) with $\tanh$ activations
in the hidden layer. In both cases, training is done using full batch gradient
descent with
$\eta=\eta_{0}=0.01,d_{\mathrm{in}}=1000,d_{\mathrm{out}}=5,\epsilon=10^{-4}$.

###  4.2 Non-linear Activations

The final extension of our work is to consider the network in Section 4.1, but
choosing nonlinear activation functions for the hidden layer. In the limit of
large $d_{h}\gg 1$, we expect the network to begin to linearize, eventually
converging to the Neural Tangent Kernel (NTK) regime [9]. 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. 5 (bottom row), we show that the dynamics of a 2 layer MLP
(1000-200-5) with $\tanh$ activations is well approximated by Eqs. 18 and 19,
empirically verifying that our predictions hold beyond the linear regimes, in
some cases.

##  5 Discussion

We have shown that grokking 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-flow limit and for large sample sizes,
were corroborated against numerical experiments and provide an excellent
description of the dynamics. In addition, preliminary evidence shows that some
of the results are applicable also beyond the linear 1-layer setting, and are
also pertinent for deeper networks and in the presence of non linearity.

Qualitatively, for linear networks with MSE loss, the training and
generalization losses are given by the squared norm of the difference between
the student and teacher weights, calculated with respect to the metric defined
by the respective covariance matrices. Training reduces both norms, and
grokking in this context simply reflects the fact that the generalization loss
lags behind the training loss. However, no qualitative change in the behavior
occurs at grokking, and consequently in this setting grokking does not imply
any transition between memorization and understanding.

It would be interesting to go beyond the gradient flow limit, and study
multilayer networks in the large learning rate regime, combining catapult/edge
of stability dynamics with grokking analysis. Additionally, studying the
effect of different optimizers on grokking could provide insights into how
algorithmic choices influence the memorization to generalization transition.
Furthermore, extending the grokking analysis to non-gaussian data or
correlated inputs could reveal how data structure and correlations affect
understanding versus memorization. Finally, in ongoing work, we consider more
realistic accuracy measures such as softmax, or cross-entropy instead of mean
squared error and connect these theoretical studies to practical deep learning
settings. Overall, understanding grokking by building upon the insights
provided by the linear estimator analysis could lead to a deeper understanding
of how artificial neural networks balance fitting the training data with
generalizing to new examples.

##  6 Acknowledgements

We thank Nadav Cohen 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.

## References

  * [1] Antoine Bodin and Nicolas Macris.  Gradient flow in the gaussian covariate model: exact solution of learning curves and multiple descent structures, 2022.
  * [2] Eric Bodin and Nicolas Macris.  Dynamics of generalization in learning with gradient descent for piecewise linear neural networks.  Advances in Neural Information Processing Systems, 34, 2021.
  * [3] Bilal Chughtai, Lawrence Chan, and Neel Nanda.  A toy model of universality: Reverse engineering how networks learn group operations, 2023.
  * [4] Andrea Crisanti and Haim Sompolinsky.  Dynamics of learning in deep linear neural networks: A mean-field approach.  Physical Review X, 8(4):041043, 2018.
  * [5] Xander Davies, Lauro Langosco, and David Krueger.  Unifying grokking and double descent.  arXiv preprint arXiv:2303.06173, 2023.
  * [6] Edgar Dobriban and Stefan Wager.  High-dimensional asymptotics of prediction: Ridge regression and classification, 2015.
  * [7] Sebastian Goldt, Marc M’ezard, Florent Krzakala, and Lenka Zdeborov’a.  Modelling the infinite width limit of neural networks with mean field theory.  AISTATS, pages 1028–1039, 2020.
  * [8] Anders Krogh and John A. Hertz.  A simple weight decay can improve generalization.  In NIPS, 1991.
  * [9] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Jascha Sohl-Dickstein, and Jeffrey Pennington.  Wide neural networks of any depth evolve as linear models under gradient descent.  Advances in neural information processing systems, 32, 2019.
  * [10] Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams.  Towards understanding grokking: An effective theory of representation learning.  Advances in Neural Information Processing Systems, 35:34651–34663, 2022.
  * [11] Ziming Liu, Eric J. Michaud, and Max Tegmark.  Omnigrok: Grokking beyond algorithmic data, 2023.
  * [12] Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mezard, and Lenka Zdeborova.  Learning curves of generic features maps for realistic datasets with a teacher-student model.  Journal of Statistical Mechanics: Theory and Experiment, 2022(11):114001, nov 2022.
  * [13] V A Marčenko and L A Pastur.  Distribution of eigenvalues for some sets of random matrices.  Mathematics of the USSR-Sbornik, 1(4):457, apr 1967.
  * [14] William Merrill, Nikolaos Tsilivis, and Aman Shukla.  A tale of two circuits: Grokking as competition of sparse and dense subnetworks, 2023.
  * [15] Beren Millidge.  Grokking ’grokking’, 2022.
  * [16] Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt.  Progress measures for grokking via mechanistic interpretability.  arXiv preprint arXiv:2301.05217, 2023.
  * [17] Pascal Jr. Tikeng Notsawo, Hattie Zhou, Mohammad Pezeshki, Irina Rish, and Guillaume Dumas.  Predicting grokking long before it happens: A look into the loss landscape of models which grok, 2023.
  * [18] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra.  Grokking: Generalization beyond overfitting on small algorithmic datasets.  arXiv preprint arXiv:2201.02177, 2022.
  * [19] Dominic Richards, Jaouad Mourtada, and Lorenzo Rosasco.  Asymptotics of ridge (less) regression under general source condition, 2021.
  * [20] Hyunjune Sebastian Seung, Haim Sompolinsky, and Naftali Tishby.  Statistical mechanics of learning from examples.  Physical review A, 45(8):6056, 1992.
  * [21] Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind.  The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon.  arXiv preprint arXiv:2206.04817, 2022.
  * [22] Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger B. Grosse.  Three mechanisms of weight decay regularization.  CoRR, abs/1810.12281, 2018.
  * [23] Bojan Žunkovič and Enej Ilievski.  Grokking phase transitions in learning local rules with gradient descent, 2022.

##  Appendix A Experimental Details

In all of our experiments, we employ a teacher-student model with shared
architecture for both teacher and student. The training data consists of a
fixed number of training samples quoted in the main text for each experiment,
drawn from a normal distribution $\mathcal{N}(0,{\bm{I}})$. All experiments
are done on MLPs using MSE loss with the 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\sim\mathcal{N}(0,1/(2d_{l-1}d_{l})$ for each layer, where $d_{l-1}$ is in
the input dimension coming from the previous layer and $d_{l}$ is the output
dimension of the current layer.

##  Appendix B Derivation of the grokking time difference

Here, we provide the full derivation for the grokking time difference
presented in Eq. 12. Our starting point is the exact solution for the training
loss in $d_{\mathrm{out}}=1$ case for $\lambda<1$, given by

$\displaystyle\mathcal{L}_{\mathrm{tr}}=e^{-4\eta_{0}(\lambda+1)t}\,_{0}\tilde{F}_{1}\left(2;16\eta_{0}^{2}t^{2}\lambda\right),$
(20)

where ${}_{0}\tilde{F}_{1}\left(a;z\right)={{}_{0}{F}_{1}}(a;z)\Gamma(a)$ is
the regularized confluent hypergeometric function. We also note the relation

$\displaystyle\frac{d\mathcal{L}_{\mathrm{gen}}}{dt}=-4\eta_{0}\mathcal{L}_{\mathrm{tr}},$
(21)

which we will use to relate training and generalization loss functions. Since
we are interested in the late time behavior, where grokking occurs, we expand
the training loss for $\eta_{0}t\gg\sqrt{\lambda}$, which is given at leading
order by

$\displaystyle\mathcal{L}_{\mathrm{tr}}\simeq\frac{e^{-4\eta_{0}\left(1-\sqrt{\lambda}\right)^{2}t}}{16\sqrt{\pi}\lambda^{3/4}(\eta_{0}t)^{3/2}}.$
(22)

Plugging in the result of Eq. 22 into Eq. 21 and integrating over time, we
find the expression for the generalization loss at late times is given by

$\displaystyle\mathcal{L}_{\mathrm{gen}}\simeq\frac{\sqrt{\eta_{0}t}e^{-4\eta_{0}\left(1-\sqrt{\lambda}\right)^{2}t}}{2\sqrt{\pi}\eta_{0}t\lambda^{3/4}}-\frac{\left(1-\sqrt{\lambda}\right)\Gamma\left(\frac{1}{2},4\eta_{0}t\left(1-\sqrt{\lambda}\right)^{2}\right)}{\sqrt{\pi}\lambda^{3/4}},$
(23)

where $\Gamma(a,z)=\int_{z}^{\infty}dte^{-t}t^{a-1}$ is the incomplete gamma
function. Expanding the result further for late times, we arrive at the result
quoted in Eq. 11. In Fig. 6, 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 at hand, we turn to the grokking time itself. As
described in the main text, we define the grokking time as the time difference
between the training and generalization accuracies reaching
$\operatorname{Erf}(\sqrt{2})\approx 95\%$, obtained when each loss satisfies
$\mathcal{L}(t^{*})=\epsilon/4$. Solving this equation for each loss
separately, in the late time limit, gives the following expressions for the
training and generalization times

$\displaystyle t^{*}_{\mathrm{tr}}$
$\displaystyle\simeq\frac{3}{8\eta_{0}\left(1-\sqrt{\lambda}\right)^{2}}\mathcal{W}\left(\frac{2\
2^{2/3}\sqrt[3]{\frac{\lambda^{3/2}}{\pi}+\frac{1}{\pi\lambda^{3/2}}-\frac{6\lambda}{\pi}+\frac{15\sqrt{\lambda}}{\pi}-\frac{6}{\pi\lambda}+\frac{15}{\pi\sqrt{\lambda}}-\frac{20}{\pi}}}{3\epsilon^{2/3}}\right),$
(24) $\displaystyle t^{*}_{\mathrm{gen}}$
$\displaystyle\simeq\frac{3}{8\eta\left(1-\sqrt{\lambda}\right)^{2}}\mathcal{W}\left(\frac{2^{5/3}\sqrt[3]{\frac{1}{\pi\lambda^{3/2}}+\frac{1}{\pi\sqrt{\lambda}}-\frac{2}{\pi\lambda}}}{3\epsilon^{2/3}}\right),$
(25)

where $\mathcal{W}(z)$ is the Lambert W function, which solves the equation
$\mathcal{W}e^{\mathcal{W}}=z$, also known as the product-log function. As the
argument of both training and generalization times are large, we can expand
the Lambert function to leading order in $z$ as $\mathcal{W}(z)\simeq\log(z)$.
Taking the difference $\Delta
t_{\mathrm{grok}}=t^{*}_{\mathrm{gen}}-t^{*}_{\mathrm{tr}}$ and expanding to
leading order in $\epsilon\ll 1$, we obtain the final expression

$\displaystyle\Delta
t_{\mathrm{grok}}=t^{*}_{\mathrm{gen}}-t^{*}_{\mathrm{tr}}\simeq\frac{\log\left(\frac{1}{1-\sqrt{\lambda}}\right)}{2\eta_{0}\left(1-\sqrt{\lambda}\right)^{2}}+\frac{3}{8\eta\left(1-\sqrt{\lambda}\right)^{2}}\log\left(1+\frac{\log\left(\left(1-\sqrt{\lambda}\right)^{4/3}\right)}{\log\left(\frac{2\left(2-2\sqrt{\lambda}\right)^{2/3}}{3\sqrt[3]{\pi}\sqrt{\lambda}{\epsilon}^{2/3}}\right)}\right),$
(26)

where the second term goes to zero as $\epsilon\to 0$, quoted in the main text
as Eq. 12.

Figure 6:  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. 22, in dashed blue, we show Eq. 23, while dotted-dashed
blue is the solution given in the main text, Eq. 11. Clearly, the asymptotic
behavior matches the exact solutions. Here,
$\eta_{0}=0.01,\lambda=0.9,d_{\mathrm{out}}=1.$

##  Appendix C Derivation for $d_{\mathrm{out}}>1$

Here, we provide additional details on the derivation of Eq. 13. The starting
point is the training and generalization loss functions, given by

$\displaystyle\mathcal{L}_{\mathrm{tr}}=\frac{1}{d_{\mathrm{out}}}\textrm{Tr}\left[D^{T}\Sigma_{\mathrm{tr}}D\right],\qquad\mathcal{L}_{\mathrm{gen}}=\frac{1}{d_{\mathrm{out}}}\textrm{Tr}\left[D^{T}\Sigma_{\mathrm{gen}}D\right]=\frac{1}{d_{\mathrm{out}}}\|D\|^{2}\
.$ (27)

where $S,T\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}$ are the
student and teacher weight matrices,
$\Sigma_{\mathrm{tr}}\equiv\frac{1}{N_{\mathrm{tr}}}\sum_{i=1}^{N_{\mathrm{tr}}}x_{i}x^{T}_{i}$
is the empirical data covariance, or Gram matrix for the training set, and we
define $D\equiv S-T$, the difference between the student and teacher matrices.
$T$ and $S$ are drawn at initialization from normal distributions
$S_{0},T\sim\mathcal{N}(0,1/({2d_{\mathrm{in}}}d_{\mathrm{out}}))$. 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 in this
instance are simply

$\displaystyle{D}_{t+1}=\left({\bm{I}}-\frac{2\eta}{d_{\mathrm{out}}}\Sigma_{\mathrm{tr}}\right)D_{t},$
(28)

where the only difference between the $d_{\mathrm{out}}=1$ case and the
equation above is the rescaled learning rate $\eta\to\eta/d_{\mathrm{out}}$
and the dimensions of $D_{t}$. Since the MP distribution is identical for each
column of $D_{t}$, the results sum up and are identical to the
$d_{\mathrm{out}}=1$ case for the losses, apart from a factor of
$1/d_{\mathrm{out}}$ and the learning rate rescaling, leading to Eq. 13.

##  Appendix D Loss calculations for Dynamics including Weight Decay

Here, we provide the derivation for Eq. 15. We begin with the definitions of
the loss function in the $d_{\mathrm{out}}=1$ case

$\mathcal{L}_{\mathrm{tr}}=D(t)^{T}\Sigma_{\mathrm{tr}}D(t),$ (29)

where $D(t)=S(t)-T$ is the difference between the student and the teacher
vectors,
$\Sigma_{\mathrm{tr}}=\frac{1}{N_{\mathrm{tr}}}\sum_{i=1}^{N_{\mathrm{tr}}}x_{i}x_{i}^{T}$
is the training covariance matrix, and $\gamma\geq 0$ is the weight decay
parameter. Using the gradient descent equation in the gradient flow limit,
$\frac{\partial D}{\partial t}=-\eta\nabla_{D}\mathcal{L}$, we obtain from Eq.
3 that

$\frac{\partial D}{\partial
t}=-2\eta\left(\Sigma_{\mathrm{tr}}+\frac{1}{2}\gamma I\right)D-\eta\gamma T.$
(30)

Multiplying by the integration factor
$e^{2\eta\left(\Sigma_{\mathrm{tr}}+\frac{1}{2}\gamma I\right)t}$ and taking
the integral, we arrive at

$D(t)+\frac{1}{2}\gamma\left(\Sigma_{\mathrm{tr}}+\frac{1}{2}\gamma
I\right)^{-1}T=e^{-2\eta\left(\Sigma_{\mathrm{tr}}+\frac{1}{2}\gamma
I\right)t}\left[D(0)+\frac{1}{2}\gamma\left(\Sigma_{\mathrm{tr}}+\frac{1}{2}\gamma
I\right)^{-1}T\right].$ (31)

We note that now the limiting value of $D(t\rightarrow\infty)$ is not zero,
but rather
$D_{\infty}=-\frac{1}{2}\gamma\left(\Sigma_{\mathrm{tr}}+\frac{1}{2}\gamma
I\right)^{-1}T$. Next, we wish to calculate
$\mathcal{L}_{\mathrm{tr}}=D(t)^{T}\Sigma_{\mathrm{tr}}D(t)$ and
$\mathcal{\tilde{L}}_{\mathrm{gen}}=D(t)^{T}\Sigma_{\mathrm{gen}}D(t)$, where
we emphasize that in both cases $D(t)$ is given by Eq. (31) and depends on
$\Sigma_{\mathrm{tr}}$. As described in the main text, it is a good
approximation to set $\Sigma_{\mathrm{gen}}$ to be the identity matrix. For
convenience, we will write both cases by
$\mathcal{L}_{\mathrm{tr}/\mathrm{gen}}=D(t)^{T}QD(t),$ where
$Q=\Sigma_{\mathrm{tr}}$ for the train and $Q={\bm{I}}$ (the identity matrix)
for the generalization.

We continue by diagonalizing $\Sigma_{\mathrm{tr}}$; we write
$M=P^{T}\Sigma_{\mathrm{tr}}P$, where $M$ is a diagonal matrix whose
eigenvalues follow the MP distribution. Hence, we obtain

$\mathcal{L}_{\mathrm{tr}/\mathrm{gen}}=\bar{D}(t)^{T}\bar{Q}\bar{D}(t),$ (32)

where $\bar{Q}=M,I$ for the train, generalization correspondingly, and
$\bar{D}(t)$ is given by

$\bar{D}(t)=e^{-2\eta\left(M+\frac{1}{2}\gamma
I\right)t}\left[\bar{D}(0)+\frac{1}{2}\gamma\left(M+\frac{1}{2}\gamma
I\right)^{-1}\bar{T}\right]-\frac{1}{2}\gamma\left(M+\frac{1}{2}\gamma
I\right)^{-1}\bar{T},$ (33)

where $\bar{D}(t)=P^{T}D(t),\bar{T}=P^{T}T$. We notice now that the expression
in Eq. (32) involves terms in the form of: $V^{T}f(M)W$ where $V,W$ are some
vectors, and $f(M)$ is some function of the diagonal MP matrix. If $V,W$ are
random vectors in a large dimension, we can approximate that

$V^{T}f(M)W=\begin{cases}0&V\neq W,\\\ |V|^{2}\int f(u)p(u)du&V=W,\end{cases}$
(34)

where $|V|$ is the norm of $V$, and $p(u)$ is the probability density function
of the MP distribution. For example, in our case we will get that
$D^{T}(0)f(M)T=-|T|^{2}\int f(u)p(u)du$ (since $D(0)=S(0)-T)$. All that is
left now is to calculate the expression in Eq. (32) explicitly, using the
approximation of Eq. (34). Doing this, at last we arrive into

$\mathcal{L}_{\mathrm{tr}/\mathrm{gen}}=d_{\mathrm{in}}\int\left(|S(0)|^{2}e^{-4\eta\left(u+\frac{1}{2}\gamma\right)t}+|T|^{2}\left(\frac{e^{-2\eta\left(u+\frac{1}{2}\gamma\right)t}u+\frac{1}{2}\gamma}{u+\frac{1}{2}\gamma}\right)^{2}\right)q_{\mathrm{\mathrm{tr}/gen}}p(u)du,$
(35)

where $q_{\mathrm{tr}}=u$ and $q_{\mathrm{gen}}=1$. By also setting the
student initialization and teacher vector norms to $|S(0)|,|T|\simeq
1/\sqrt{2d_{\mathrm{in}}}$ (as done in the main text), we finally get

$\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}=\frac{1}{2}\int\left(e^{-4\eta\left(u+\frac{1}{2}\gamma\right)t}+\left(\frac{e^{-2\eta\left(u+\frac{1}{2}\gamma\right)t}u+\frac{1}{2}\gamma}{u+\frac{1}{2}\gamma}\right)^{2}\right)q_{\mathrm{\mathrm{tr}/gen}}p(u)du.$
(36)

##  Appendix E Derivation for the 2-layer network

Here, we provide supplementary details on the derivation of Eq. 17. We
consider the addition of a single hidden linear layer, where the teacher
network function is $f(x)=(T^{(1)})^{T}(T^{(0)})^{T}x$, where
$T^{(0)}\in\mathbb{R}^{d_{\mathrm{in}}\times d_{h}}$,
$T^{(1)}\in\mathbb{R}^{d_{h}\times d_{\mathrm{out}}}$ and $d_{h}$ is the width
of the hidden layer. Similarly, the student network is defined by two matrices
$S^{(0)},S^{(1)}$. The empirical training loss over a sample set
$\\{x_{i}\\}_{i=1}^{N}$ reads

$\displaystyle\mathcal{L}_{\mathrm{tr}}$
$\displaystyle=\frac{1}{N_{\mathrm{tr}}d_{\mathrm{out}}}\sum_{i=1}^{N_{\mathrm{tr}}}\left((S^{(1)})^{T}(S^{(0)})^{T}x_{i}-(T^{(1)}))^{T}(T^{(0)}))^{T}x_{i}\right)^{2}\
.$ (37)

In this setup the weights are drawn at initialization from normal
distributions
$S^{(0)}_{0},T^{(0)}\sim\mathcal{N}(0,1/({2d_{\mathrm{in}}}d_{h}))$ ,
$S_{0}^{(1)},T^{(1)}\sim\mathcal{N}(0,1/({2d_{\mathrm{out}}}d_{h}))$. Next, we
define $T=T^{(0)}T^{(1)}\in\mathbb{R}^{d_{\mathrm{in}}\times
d_{\mathrm{out}}}$ and derive the gradient flow equations for the system

$\displaystyle\dot{S}_{t}^{(0)}=-\frac{2\eta_{0}}{d_{\mathrm{out}}}\Sigma_{\mathrm{tr}}\left(S^{(0)}_{t}S^{(1)}_{t}-T\right)(S^{(1)}_{t})^{T},\quad\dot{S}_{t}^{(1)}=-\frac{2\eta_{0}}{d_{\mathrm{out}}}(S^{(0)}_{t})^{T}\Sigma_{\mathrm{tr}}\left(S^{(0)}_{t}S^{(1)}_{t}-T\right).$
(38)

defining $D_{t}=S_{t}^{(0)}S_{t}^{(1)}-T$, and noting that
$\dot{D}_{t}=S_{t}^{(0)}\dot{S}_{t}^{(1)}+\dot{S}_{t}^{(0)}S_{t}^{(1)}$, we
arrive at the equations quoted in the main text

$\displaystyle\dot{D}_{t}=-2\eta_{0}\frac{h}{d_{\mathrm{out}}^{2}}\Sigma_{\mathrm{tr}}D,\qquad\dot{h}_{t}=-8\eta_{0}(T+D)^{T}\Sigma_{\mathrm{tr}}D.$
(39)

Here, $h=Tr[H]/2=\|S^{(0)}\|^{2}/2+\|S^{(1)}\|^{2}/2$, where
$H=\nabla_{\theta}^{T}\nabla_{\theta}\mathcal{L}_{\mathrm{tr}}$ is the Hessian
matrix and $\theta\equiv\\{S^{(0)},S^{(1)}\\}$. Although Eq. 18 describes a
set of coupled equations, we note that the solution for $h_{t}$ can be
simplified when considering

the limit of small $\eta_{0}\ll 1$, as we may ignore the time evolution and
consider the trace (or kernel) as fixed to its initialization value, which is
$h_{0}\simeq 1/2$ for $d_{h}\gg d_{\mathrm{out}}$. In that case the loss
solutions are a simple modification to the ones given in the previous
sections, with the replacement $\eta_{0}\to\eta_{0}/(2d_{\mathrm{out}}^{2})$.
Subsequently, the training/generalization performance metrics are

$\displaystyle\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}^{\mathrm{2-layer}}$
$\displaystyle=\|D_{0}\|^{2}\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}^{\mathrm{1-layer}}\left(\frac{\eta_{0}}{2d_{\mathrm{out}}^{2}},\lambda,t\right),\quad\mathcal{A}_{\mathrm{\mathrm{tr}/gen}}=1-\frac{\Gamma\left(\frac{d_{\mathrm{out}}}{2},\frac{d_{\mathrm{out}}\epsilon}{2\mathcal{L}_{\mathrm{\mathrm{tr}/gen}}}\right)}{\Gamma\left(\frac{d_{\mathrm{out}}}{2}\right)}.$
(40)