Predicting Grokking Long Before it Happens:
A look into the loss landscape of models which grok

Pascal Jr. Tikeng Notsawoyyy,sch

Hattie Zhouyyy,sch
Mohammad Pezeshkimeta

Irina Rishyyy,sch
Guillaume Dumassch,comp

yyyMila, Montréal, Quebec, Canada
compCHU Sainte-Justine Research Center, Montréal, Quebec, Canada
schUniversité de Montréal, Montréal, Quebec, Canada
metaMeta AI Research

Pascal Tikeng<EMAIL_ADDRESS>
Machine Learning, ICML

This paper focuses on predicting the occurrence of grokking in neural networks, a phenomenon in which perfect generalization emerges long after signs of overfitting or memorization are observed. It has been reported that grokking can only be observed with certain hyper-parameters. This makes it critical to identify the parameters that lead to grokking. However, since grokking occurs after a large number of epochs, searching for the hyper-parameters that lead to it is time-consuming.  In this paper, we propose a low-cost method to predict grokking without training for a large number of epochs. In essence, by studying the learning curve of the first few epochs, we show that one can predict whether grokking will occur later on. Specifically, if certain oscillations occur in the early epochs, one can expect grokking to occur if the model is trained for a much longer period of time. We propose using the spectral signature of a learning curve derived by applying the Fourier transform to quantify the amplitude of low-frequency components to detect the presence of such oscillations. We also present additional experiments to explain the cause of these oscillations and characterize the loss landscape.

§ INTRODUCTION

Despite the recent growth of theoretical studies and empirical successes of neural networks [Graves et al., 2013, He et al., 2015, Krizhevsky et al., 2012, Silver et al., 2016], understanding why such networks

find generalizable solutions
in over-parameterized regimes, where the number of learnable parameters is much larger than the number of training samples, remains an open question. Indeed, one of the major challenges of deep learning is that, in practice, neural networks are highly overparameterized [Allen-Zhu et al., 2019, Zhang et al., 2017], and in some cases, it is observed that overparameterization empirically improves optimization and generalization, which seems to contradict traditional learning theory. For example, Livni et al., 2014 observed that in synthetic data generated from a target network, the learned network converges faster when it has more parameters than the target network. Arora et al., 2018 also found that, in practice, trained overparameterized networks can often be compressed into simpler networks with far fewer parameters, without affecting their generalizability. Nakkiran et al., 2020 experiment a double-descent phenomenon where, as the model size increases, performance first gets worse and then gets better.  Recently, Power et al., 2022 have shown through a phenomenon they named grokking that long after severe overfitting,

validation accuracy sometimes suddenly begins to increase from chance level to perfect generalization, and that the amount of optimization required for this generalization increases rapidly as the size of the data set decreases.

The grokking phenomenon opens the way to new studies concerning the structure of the minimum found by Stochastic Gradient Descent (SGD), and how networks behave in the neighbourhood of SGD training convergence. Indeed, neural activity is often characterized by an exploratory early phase of rapid learning with a rapid and sometimes abrupt decrease in the loss function [Nakkiran et al., 2020, Feng & Tu, 2021]. This phase is followed by a second phase, often called the diffusion phase [Shwartz-Ziv & Tishby, 2017], when the learning error reaches its minimum value and the global loss decreases again, but much more slowly and gradually. This second phase, often chaotic [Herrmann et al., 2022, Cohen et al., 2021, Thilak et al., 2022], is characteristic of the structure of the minimum found by the optimization algorithm [Goodfellow & Vinyals, 2015, Im et al., 2016, Smith & Topin, 2017, Keskar et al., 2017, Li et al., 2018, Jastrzębski et al., 2018, Feng & Tu, 2021].
Optimization hyper-parameters and initialization strongly affect training dynamics during the above phases, the convergence to a specific region in parameter space, and the geometries and generalization properties of solutions found by SGD [Keskar et al., 2017, Jastrzębski et al., 2018]. For example, Keskar et al., 2017 empirically showed that a larger batch size correlates with sharper or isolated  minima and worse generalization performance, while small batch methods correlates with flater minima. Jastrzębski et al., 2018 find that the critical control parameter for SGD is not the batch size alone, but the ratio of the learning rate and the batch size, and that higher values for this ratio result in convergence to wider minima.

Recent work has shown that grokking is observed only with a certain range of hyperparameters [Power et al., 2022, Liu et al., 2022]. Others have studied microscopic phenomena that coincide or come in tandem with delayed generalization, such as the emergence of structure in embedding space [Liu et al., 2022] and the slingshot effect [Thilak et al., 2022].

But observing grokking often requires training the model for a very long time, making it difficult to construct a phase diagram of generalization covering all the hyperparameters.
In this work, we propose a low-cost method that can predict grokking long before it occurs. Our main contributions are:

* We study the learning curves of a transformer network [Vaswani et al., 2017] trained on arithmetic data in settings with and without grokking (section <ref>).

* We propose spectral signature to quantify the oscillations of the loss in the early phases of training, and show that empirically it can be used to predict a potential generalization

(section <ref>).

To understand the origin of the oscillations and the shape of the minimizer, we analyze the model's loss landscape along the training trajectory, and present evidence linking the shape of this landscape to the different phenomena related to grokking such as the slingshot mechanism [Thilak et al., 2022] (section <ref>).

§ PRELIMINARIES

§.§ Notations

The index $t$ is used to characterize the training steps ($\theta_t$, etc.).  $L_t = L(\theta_t)$ denotes the loss at $\theta_t$ (the parameter update at time $t$ given the optimization algorithm of choice), $G_t = \nabla L(\theta_t)$ the gradient of the loss function at  $\theta_t$  and $\mathcal{H}_t = \nabla^2 L(\theta_t)$ the local Hessian matrix of the loss at $\theta_t$ ($\theta(t)$, $L(t)$, $G(t)$ and $\mathcal{H}(t)$ in the continuous case). We define $\dot{x} =  \partial x (t) /  \partial t$ for a time function $x = x(t)$.

Task definition and training scheme

Let $\circ$ be a binary mathematical operator (such as addition or multiplication), and let $p$ and $q$ be two strictly positive integers. The task is to predict $(a \circ b) \ mod \ q$ for any pair of numbers $(a, b) \in [p]^2$, with $[p] = \{0, \dots, p-1\}$. The dataset $\mathcal{D}$ that we can thus constitute has a size of $p(p+1)/2$ if $\circ$ is symmetric (and we consider $a \circ b$ and $b \circ a$ as the same operation), and $p^2$ otherwise. $\mathcal{D}$ is randomly partitioned into two disjoint and non-empty sets $\mathcal{D}_{train}$ and $\mathcal{D}_{val}$, the training and the validation dataset respectively. The training data fraction $r = |\mathcal{D}_{train}| / |\mathcal{D}|$ is a  hyperparameter.

This problem can be solved with an auto-regressive approach. For $a, b \in [p]$, let $s = s_1 \dots s_5 = \langle a \rangle \langle \circ \rangle \langle b \rangle \langle = \rangle \langle (a \circ b) \ mod \ q \rangle$ where $s_i = \langle x \rangle$

stands for the token corresponding to the element $x$.
The training is performed by maximizing the likelihood under the direct autoregressive factorization, and the loss (as well as the accuracy) is calculated only on the answer part $s_5$ of the equation :

\max_{\theta} \log \ p_{\theta} (s_5 | s_{<5})
%log \mbox{ } \frac {exp(W_{ s_5 }^T h_{\theta}(s_{<5}))}
%{\sum_{x \in [q]} exp(W_{\langle x \rangle}^T h_{\theta}(s_{<5}))}
\log \mbox{ } \frac {\exp(h_{\theta}(s_{<5})^T e(s_5))}
{\sum_{x \in [q]} \exp(h_{\theta}(s_{<5})^T e(\langle x \rangle ))}

where $h_{\theta}(s_{<5}) \in \mathbb{R}^{m}$ is a context representation produced by neural models parameterized by $\theta$ for $s_{<5} = s_1 \dots s_4$;

and $e(s_i)$ the trainable embedding  vector of $s_i$ (the embedding weights are used as weights for the output linear classification layer).

We focus on the following binary operations, for a prime number $p = 97$, and $q = p$ : modular addition ($+$), $a \circ b = a + b \ (mod \ p)$ for $a, b \in [p]$;
multiplication in the permutation group $S_5$, $a \circ b = a \cdot b$ for $a, b \in S_5$.

For all experiments, we used a transformer with 2 layers, 128 embeddings, and 4 attention heads, with a total of roughly $4 \cdot 10^5$ non-embedding parameters. We use the following hyperparameters unless stated otherwise: AdamW optimizer with a learning rate of $10^{-4}$, weight decay of $1$ and ($\beta_1, \beta_2) = (0.9, 0.98)$, linear learning rate warm up over the first $10$ updates, and minibatch size of $|\mathcal{D}_{train}|$.

Grokking and Non-grokking

The phenomenon of grokking is characterized by three phases in order: 1) the initial learning phase where both training and validation performance are poor, 2) the second phase where training performance is near perfect while validation performance is low, 3) and the generalization phase where validation performance improves to match training performance. The grokking phenomenon is marked by a long memorization phase (figure  <ref>). More specifically, let $t_1$ be the transition step from phase 1 to phase 2 (when the training accuracy reaches a value strictly greater than 0% for the first time), $t_2$ the step at which the training data is completely memorized (when the training accuracy reaches 100% for the first time), $t_3$ the transition step from memorization to comprehension (when the validation accuracy reaches a value strictly greater than 0% for the first time), and $t_4$ the step at which the model generalizes (when the validation accuracy reaches 100% for the first time). It is the delay between memorization and generalization $t_4 - t_2$

that allow, in general, to name the learning process grokking or simple generalization. But in our case, for the sake of simplicity, we will characterize all situations with generalization as grokking.

Oscillation in training and validation accuracies. We train on modular addition with $r=0.5$.

The left curve shows a case where the model did not grok after 10k steps of training. The right curve shows the generalization after overfitting. Training accuracy becomes close to perfect at $t_2 < 300$ optimization steps, but it takes close to $t_4 \approx 7k$ steps for validation accuracy to reach that level.

Although it is easy to identify grokking, it is very difficult to give a formal definition of its opposite, since in practice we cannot train a model for an infinite number of steps. We proceeded empirically by training the models for a large range of hyperparameters, and we fit a function that predicts $t_4$, the generalization step, for each training data fraction $r$. That is, if for other hyper-parameters, and for a given training data fraction $r$, we train a model for $t_4(r)+\epsilon$ (we used $\epsilon = 1k$) steps and have no generalization, we can stop the training. In general, more data leads to faster grokking (i.e. $t_2(r)$ and $t_4(r)$ are decreasing functions of $r$). Empirically, $t_4$ follows a power law of the form $t_4(r) = a r ^{-\gamma} + b$. This was first predicted by Žunkovič & Ilievski, 2022.  We give more details in the appendix (<ref>).

§ PREDICTING GROKKING

Oscillation in training and validation accuracies. We train on modular addition with $r=0.5$. The red curves show training accuracy and the green ones show validation accuracy. The first and last curves show a case where the model did not grok after 10k steps of training.

A starting observation is that the learning curves of models that grok exhibit oscillatory behaviors

(Figure <ref>).

A related phenomenon, the slingshot effect,  was observed and named recently by Thilak et al., 2022. Cyclic transitions between a stable and unstable training regime characterize this phenomenon.

Thilak et al., 2022 characterized this phenomenon as the complete cycle starting with the norm growth phase and ending with the norm plateau phase, and found that it is ubiquitous and can be easily replicated in multiple scenarios, encompassing a variety of models and data sets. They observe that slingshots and grokking tend to come in tandem, that is grokking almost exclusively happens at the onset of slingshots, and is absent without it.
As pointed by Thilak et al., 2022, this type of transition is reminiscent of the catapult phenomenon [Lewkowycz et al., 2020], where the loss initially increases and begins to decrease once the model catapults to a region of lower curvature early in the training, for a sufficiently large training step.

Based on this observation, we conjecture that the spectral signature of training loss in early epochs can hint us about the existence of an upcoming grokking. We first try to comparatively quantify the oscillations in the training loss when the model groks and when it does not grok, the ideal being to stop the training if the model does not seem to be able to grok, in order to save computational resources.
From the gradient flow equation $\dot{\theta} = - G(t)$, it holds that $\dot{L} \approx - \| G(t) \|^2$
$\ddot{L} \approx 2 G(t)^T H(t) G(t) = 2 \sum_{i} \lambda_i(t) \langle G(t), v_i(t) \rangle^2
$\{ \lambda_i(t) \}_i$ the spectrum of $H(t)$, and $\{ v_i(t) \}_i$ the associated eigenvectors. From this, it becomes clear that the evolution of $L$ over time depends on the norm of the gradient, and how fast it changes depends on the curvature of its landscape. Any signal that can be represented as a variable that varies in time has a corresponding frequency spectrum.

We considered the training loss $L$ over the training steps (and in the early stages of training) as signals and analyzed its spectral signature. By spectral signature of the loss, we mean any measure or set of measures that can quantify the oscillations in the loss, such as the spectral energy or the Hjorth parameters [Hjorth, 1970] - activity, mobility and complexity. The Hjorth activity, which is the variance of the signal in the time domain, is equal to the spectral energy if the signal $L$ has zero mean. The latter condition is obtained

by passing $L$ through a sufficiently low-pass filter in the frequency domain, which removes its non-oscillatory components
[A non-oscillatory component of a signal is any component that does not vary rapidly with time, such as a constant value, a linear trend, or a smooth curve. These components are generally considered to be low-frequency signals, and can be removed by passing the signal through a low-pass filter.
that are not necessary for the quantification of the oscillations. In this case, the Hjorth parameters are directly related to the loss landscape as illustrated below.

Let $\mathcal{F}(L)$ denote the Fourier transform of $L(t)$ and $m_n(L) = \int \omega^n  |\mathcal{F}(L)(\omega)|^2 d \omega$, the $n^{th}$ moment of $\mathcal{F}^2 (L)$,

with $|\mathcal{F}(L)(\omega)|^2$ the energy spectral density present in the

pulse $\omega$.
The Hjorth activity represents the signal power, the surface of the power spectrum in the frequency domain. It is given by $m_0(L)$, which is equal to $\int L^2 (t) dt$ by the parseval's theorem [
The Hjorth mobility is the mean frequency or the proportion of standard deviation of the power spectrum, and is given by $\sqrt{m_2 (L) / m_0(L)}$ with
$m_2(L) = \int  | \omega \mathcal{F}(L)(\omega)|^2 d \omega
= \int \dot{L}^2 (t) dt \approx  m_0 (\dot{L})$ the activity of the gradient norm $\| G(t) \|^2$. In a similar way,
The Hjorth complexity, which indicates how the shape of a signal is similar to a pure sine wave, is given by $\sqrt{m_4 (L) / m_2(L)}$ with
$m_4(L) = \int  | \omega^2 \mathcal{F}(L)(\omega)|^2 d \omega = \int \ddot{L}^2 (t) dt \approx m_0 (\ddot{L})$ the activity of the hessian spectrum.

]. Figure <ref> shows a similarity between the oscillation in the training loss in the early phases of training and the validation accuracy for $r=0.5$ (more results in the appendix, <ref>), suggesting that the spectral signature can serve as a proxy to upcoming grokking phenomenon. The generalization is most observed for small learning rate and small weight decay. Although large learning rates have the effect of increasing the oscillations, this does not directly result in grokking and is not necessarily visible in the early steps, but more near the basin of attraction of the minimum. Importantly, the spectral signature of the loss is not an explicit capacity measure so either positive or negative correlation with generalization could potentially be informative. Our observations is related to the empirical findings of Jiang et al., 2019. They investigate more than 40 complexity measures derived from both theoretical and empirical studies and train a variety of models by systematically varying commonly used hyperparameters. Their results suggest that the difficulty of optimization during the initial phase of the optimization benefits the final generalization, but the evolution of the loss when it reaches a certain value is not correlated to the generalization of the final solution.

The first figure (top) represents the validation accuracy (%) at the end of the training ($10k$ steps), and the second figure (bottom) represents the spectral energy (activity) in the training loss for the first 400 training steps ($r = 0.5$). On the x-axis we have the weight decay strength, and on the y-axis we have the learning rate. A similarity is observed between the oscillation patterns in the training loss during the initial stages of training and the validation accuracy. This suggests that the spectral signature can be used as an indicator or proxy for the upcoming grokking phenomenon.
The highest degree of generalization is typically observed when using small learning rates and small weight decay. While large learning rates may increase oscillations, this does not directly lead to grokking and is not necessarily evident in the early stages of training. Instead, such effects become more noticeable near the basin of attraction of the minimum.

Periodogram of the training loss for the first 400 training steps over 20k in total ($r=0.5$). This y-axis is the power spectral
density computed from the time-frequency decomposition, with fast Fourier transforms for convolutions and with Morlet
wavelets. The first figure uses a frequency band from 0 to 2000
for 2000 sampling epochs
while the second is a zoom of the first one at low frequencies. The title in purple on each figure is the weight decay strength, in $\{0.0, 0.4, 0.9, 1.3, 1.8, 2.5, 3.0, 3.5\}$. The labels are the training rate, in $\{0, 1, 2, 3, 4, 5 \} \times 10^{-2}$, with the spectral energy between brackets followed by the maximum validation accuracy (%) achieved during training by the model. On the x-axis we have the weight decay strength, and on the y-axis we have the learning rate.

§ GROKKING LOSS LANDSCAPE

What can be a possible explanation of the oscillations observed in the early phases and the slingshot phenomenon? Many hypotheses can in fact be put forward to reconcile oscillations with delayed generalization. Does the model, during the confusion, memorization and comprehension phases, oscillate around a local minimum, cross a very flat region, or circumvent a large obstacle?

One initial intuitive explanation for delayed generalization could be that the model gets stuck in local solutions during the memorization phase. The ease of escaping these local solutions depends on factors such as initialization or hyperparameters like the amount of training data, and the model achieves grokking when it successfully breaks free from the basin of attraction of such solutions.
This explanation is similar to the working hypothesis of Dziugaite & Roy, 2017 that SGD finds good solutions only if they are surrounded by a relatively large volume of solutions that are nearly as good.

A second attempt to explain grokking landscape is that the model crosses an ill-conditioned surface, potentially a valley with almost no curvature in the majority of directions and very high curvature in some directions.

This results in a weak progression in the directions of low curvature, and a lot of back and forth in the directions of high curvature.

A third attempt is to consider the first two hypotheses together. That is, the model goes through several ill-conditioned local minima during training.

We present in the following the loss landscape of grokking, a discussion on how they are linked to our hypothesis.

Since neural loss functions live in a very high-dimensional space, visualizations are only possible using low-dimensional 1D (line) or 2D (surface) plots. In this work, we consider the approach of Li et al., 2018. Let $\theta$ be the point near which we want to observe the loss landscape.  We plot the loss and the accuracy as a function of $A \subseteq \mathbb{R}$, $f(\alpha) = L(\theta + \alpha \vec{\delta})$, where $\vec{\delta}$ is a direction vector carefully chosen in $\vec{\Theta}$ [
In 2D, the loss is plot as a function of $A \times B \subseteq \mathbb{R}^2$, $f(\alpha, \beta) = L(\theta + \alpha \vec{\delta} + \beta \vec{\eta})$, where $\vec{\delta}$ and $\vec{\eta}$ are two carefully chosen direction vectors in $\vec{\Theta}$.

$\vec{\delta}$ and $\vec{\eta}$ can be randomly chosen or defined by $\vec{\delta} = \theta^{'} - \theta$ and $\vec{\eta} = \theta^{''} - \theta$, with $\theta^{'}$ and $\theta^{''}$ another points whose choice will be specified.
]. Due to scale invariance in network weights, this approach may fail to capture the intrinsic geometry of loss surfaces. To remove this scaling effect, we plot the loss functions using an adaptation of the filter-wise normalized directions [Li et al., 2018], that is
w_k(\vec{\delta}) \leftarrow \frac{w_k(\vec{\delta})}{\| w_k(\vec{\delta}) \|} \| w_k(\theta) \|
\text{ and }
b_k(\vec{\delta}) \leftarrow \frac{b_k(\vec{\delta})}{| b_k(\vec{\delta}) |} | b_k(\theta) |
$ for each weight $W = [\dots, w_k, \dots]^T$ ($w_k$ is a vector) and bias $b = [\dots, b_k, \dots]$ ($b_k$ is a scalar) in each layer of $\theta$ and $\vec{\delta}$ [
The same as for $\vec{\eta}$ when applied
]. This loss-landscape visualization approach, although simple, has the advantage of allowing to visualize the potential local convexity of the loss

in the chosen direction [
A function $f: \mathbb{R}^n \to \mathbb{R}$ is convex if and only if $
g_{(x, y)} \colon
\alpha \in [0, 1] \mapsto
g_{(x, y)}(\alpha) =
f\big(x + \alpha (y-x)\big)$ is convex for all $x, y \in \mathbb{R}^n$.

For $r = 0.3$, figure <ref> shows the 1D projection of the grokking loss surface for a single epoch of training (just after the grokking step), while figures <ref> ($\vec{\delta}_t \propto \theta^* - \theta_t$), <ref> ($\vec{\delta}_t \propto \theta_{t+1} - \theta_t$), <ref> ($\vec{\delta}_t \propto \theta_0 - \theta_t$) and <ref> (random $\vec{\delta}_t$) show it for different training epochs (more experiments in the Appendix <ref>). We can see that the 1-D subspace from initial to final parameters and from one minimizer to another contains many difficult and exotic structures.

[multiplication ($S_5$)]
[addition ($+$)]
1D projection of the grokking loss and accuracy surface ($r=0.3$).
The x-axis is $\alpha \in [-3, 3]$, and the y-axis are the loss (left axis) and accuracy (right axis) at $\theta^* + \alpha \delta$ with $\delta \propto \theta_0 - \theta^*$ (with an adaptation of the  filter-wise normalization [Li et al., 2018]), where $\theta_0$ is the initial parameter and $\theta^*$ the parameter just after the model has grokked (a) and before grokking (b).
We can see that the 1-D subspace from initialization to grokking contains many difficult and exotic structures. In (a), we have two minimizers of the training loss (solid line), but only one of them also minimizes the validation loss (dashed line).

in 2, 100, 247,
400, 600, 800, 987, 1187, 1587

Until the end of the training,

$\vec{\delta}_t = \theta^* - \theta_t$
($r = 0.8$).
Here, the structure is more exotic.
It should be noted here that while $\alpha=0$ corresponds to the point around which the loss is observed ($\theta_t$), $\alpha=1$ does not necessarily correspond to the optimum $\theta^*$, since there is normalization. Around $t>200$, two minimizers of the training loss emerge at $\alpha=0$ and $\alpha=1$, and are separated from each other by a mountain. Only one of the two minimizers also minimizes the validation loss. Towards the end of the training ($t>900$, when the model generalizes), the two attractors annihilate into a single (apparent flat) global optimum.

[Loss surface : $f_t(\alpha) = Loss(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

[Accuracy surface : $f_t(\alpha) = Acc(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

1D projection of the grokking loss and accuracy surface. This corresponds to figure <ref>.a, but for several training epochs ($r=0.3$). The direction $\vec{\delta}_t$ used for each training epoch $t$ is the unit vector of  $\theta^* - \theta_t$ (filter-wise normalization [Li et al., 2018]), the direction from the parameter at epoch $t$ to the minimum.

Here, the structure is more exotic. We can clearly see two minimizers of the training loss, but only one minimizes the validation loss: during memorization, the model is in this local minimum, and it achieves
grokking when it successfully breaks free from this local solution.

[Loss surface : $f_t(\alpha) = Loss(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

[Accuracy surface : $f_t(\alpha) = Acc(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

Similar to figure <ref>,
but with $r=0.8$. What changes between the two figures is the time it takes for the model to reach the global minimum. This time decreases as $\Theta \left( 1/ r ^{\gamma} \right)$ with $r$, $\gamma > 0$

(section <ref>).

in 0, 2, 100, 159, 160,
200, 249, 349, 400, 600,
800, 987, 1187, 1968, 1971

Epoch to Epoch ($r = 0.8$), $\vec{\delta}_t = \theta_{t+1} - \theta_t$.
This allows a closer look at the behaviour of the landscape in the direction of the gradient, at each training step $t$. For some $t$, the loss, although disturbed, remains locally flat (up to a resolution used for visualization). These are the instants during which the weight progresses along the directions of least curvature, resulting in a stagnation of the loss. These points are usually followed by a slingshot.

[Loss surface : $f_t(\alpha) = Loss(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

[Accuracy surface : $f_t(\alpha) = Acc(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

1D projection of the grokking loss/accuracy surface for epoch to epoch, $\vec{\delta}_t \propto \theta_{t+1} - \theta_t \propto - \nabla L (\theta_t)$.
This allows a closer look at the behaviour of the landscape in the (local) direction of the gradient, at each training step $t$. For some $t$, the loss, although disturbed, remains locally flat (up to a resolution used for visualization). These are the instants during which the weight progresses along the directions of least curvature, resulting in a stagnation of the loss. These points are usually followed by a slingshot.

To measure the level of curvature of the loss function, we compute the maximum (${\lambda}_{max}$) and minimum (${\lambda}_{min}$) eigenvalue of its Hessian (more details in the Appendix <ref>). We observe that there is no significant negative curvature in the trajectory. The curvature remains generally positive and is greatly disturbed at the slingshot points (${\lambda}_{min}$ remains in general close to 0 while ${\lambda}_{max}$ is large, but during slingshot ${\lambda}_{min}$ becomes negative). Our loss is ${\lambda}_{max}$-smooth, i.e. its gradient is ${\lambda}_{max}$-Lipschitz. It is known from optimization literature that a function with bounded Hessian eigen-values has the gradient that tend to decay when the parameter gets closer to the minimum, in constrast to non-smooth one that generaly have abrupt bends at the minimum, which cause significant oscillations for gradient descent [Bubeck et al., 2015]. We also observed that more than 98% of the total variance in the parameter space occurs in the first 2 PCA modes much smaller than the total number of weights, suggesting that the optimization dynamics are embedded in a low-dimensional space [Li et al., 2018, Feng & Tu, 2021]. Moreover, the model remains in a lazy training regime [Chizat et al., 2019, Berner et al., 2021] most of the time, as the measure of cosine distance between the model weights from one training step to the next remains almost constant, except at the slingshot location.

These observations seem to support the hypothesis that the model crosses a perturbed valley of bad solutions and experiences disturbances along the direction of the valley. The valley here is just a surface with a large curvature in most directions and no (low) curvature in the rest of the directions. Hence the learning activity of SGD is governed by these directions of high curvature.  When the iterates fall in the valley, we are at the minimum for the training objective (figure <ref>), so that the model can memorize the training data. This local minimum in the valley directions is flat, but in the high curvature directions, is sharp. With the accelerated or adaptative gradient methods, the pathological curvature is no longer a problem, but the local minimum and saddle points are.
This interpretation of grokking was the working hypothesis of Dziugaite & Roy, 2017: SGD finds good solutions only if they are surrounded by a relatively large volume of solutions that are nearly as good, as we saw above with grokking optimum that is surrounded by many local minima along the principal directions of curvature. Their work focuses on showing that, under realistic hypotheses, SGD performs implicit regularization or tends to find solutions that possess some particular structural property that we already know to be connected to generalization, like widder minima, that are less difficult to reach by SGD than sharper one as they have a large basin of attraction. Moreover, the loss function here does not seem to satisfy the strict saddle property ($\lambda_{min} < 0$), which guarantees for gradient descent to not converge to saddle points on continuously differentiable functions when using a step size smaller than $1/\lambda_{max}$ [Lee et al., 2016]. The dynamic, therefore, has components in the unstable subspace of the Hessian (a sub-space spanned by the dimensions of non-negative eigenvalues), and if $\theta_0$ is chosen out of this space (which happens with zero probability for a random initialization), there is hope that we will quickly converge to a global solution. The spectral signature reflects the dynamics in and out of this unstable sub-space. The strict saddle property is indeed verified only when there are strong oscillations.

The Slingshot mechanism

Above, we derive $\dot{L}(t) = - \| G(t) \|^2$ for a small enough step size. This relaxation property of gradient descent reflects the fact that the loss function cannot increase. However, we lose this property with certain accelerated or adaptative methods like Adam, or on ill-conditioned problems as we saw above with $L_{t+1} > L_t$ for $\epsilon_t \lambda_t(i) > 2$. In the case of grokking, the loss, even when it becomes zero, presents a sudden growth, followed by a decay, and this in a periodic way.
Thilak et al., 2022 have recently shown that this phenomenon is general to the optimization of deep neural networks. Between two slingshots, the gradient is almost zero, all eigenvalues of the Hessian are nonnegative, and there is one direction of very large curvature that dominates the others and the weights update, resulting in a small change in weight from one training step to another. The model thus seems to traverse a flat valley.

To reduce this phenomenon (i.e. the amplitude of the spikes), we clipped the gradient during training.  This had the effect of slowing down generalization, but not preventing it.

\dot{\theta} = \left\{
\begin{array}{ll}
- \eta G(t) / \| G(t) \| & \mbox{if } \| G(t) \| \ge \eta \\
- G(t) & \mbox{otherwise.}
\end{array}
\right. (\eta > 0)

\Rightarrow \dot{L} = \left\{
\begin{array}{ll}
- \eta \| G(t) \| & \mbox{if } \| G(t) \| \ge \eta \\
- \| G(t) \|^2 & \mbox{otherwise.}
\end{array}
\right. (\eta > 0)

The spikes are reduced as $\eta \rightarrow 0$, but remain visible as $\theta \rightarrow \theta^*$.  By reducing the learning rate extremely, the spikes also become less visible, but we also pay the same cost because the model takes more steps to generalize.

Effect of gradient clipping on slingshot

TODO : (1) figure (2) how this affect the signature

§ RELATED WORKS

Power et al., 2022 are the first to have studied and named the grokking phenomenon. They train a decoder-only causal transformer [Vaswani et al., 2017] to predict the response $c$ of a binary operation of the form $a \circ b=c$, where $a$, $b$, $c$ are discrete symbols without internal structure, and $\circ$ a binary operation (addition, composition of permutations, bivariate polynomials, etc). Liu et al., 2022 studied grokking in simple toy models, and embedding layer followed by a multilayer perceptron, for classification and regression. The authors empirically show that generalization on algorithmic datasets coincides with the emergence of structure in embeddings, that the quality of representation predicts generalization, and they defined the notion of representation quality in the toy setting and shown that it predicts generalization by developing an effective theory to describe the learning dynamics of the representations in the same toy setting. They also illustrate phase diagrams including memorization, comprehension, confusion and grokking (similar to comprehension, but with a very large number of training steps between memorization and generalization) as a function of hyperparameters. The taxonomy used throughout this paper is identical to theirs, with the difference that we refer in our work to the phases along the training trajectory. We examined in more detail the loss landscape during these different training phases, but we do not analyze the embedding structure.
Thilak et al., 2022's slingshot mechanism is quite correlated with our hypotheses about oscillations. Their work focuses on characterizing periodic oscillations observed in the different phases of model training, and they were the first to empirically observe that grokking almost exclusively happens at the onset of slingshots, and is absent without it. Our work reveal that these oscillations are caused at the level of the landscape by the escape of sharp mountains surrounding local optima. Moreover, the model weights change abruptly during slingshot, and slightly outside the slingshot. Other works have also been interested in grokking recently, such as the mechanistic interpretability of grokking [Nanda & Lieberum, 2022] and grokking as a phase transition [Žunkovič & Ilievski, 2022].

Neural networks loss landscape

The visualization methods use in this paper was first introduced by Goodfellow & Vinyals, 2015, and is still one the most use method in date to study the effect of loss landscapes on generalization [Goodfellow & Vinyals, 2015, Im et al., 2016, Smith & Topin, 2017, Keskar et al., 2017, Li et al., 2018, Feng & Tu, 2021]. Li et al., 2018 observe that, when networks become sufficiently deep, neural loss landscapes quickly transition from being nearly convex to being highly chaotic, and that this transition coincides with a dramatic drop in generalization error, and ultimately to a lack of trainability. Feng & Tu, 2021 investigate the connection between the the loss function landscape and the stochastic gradient descent (SGD) learning dynamics. They show that around a solution found by SGD, the loss function landscape can be characterized by its flatness in each PCA direction. As Li et al., 2018, they show that optimization trajectories lie in an extremely low dimensional space.

Flatness, Sharpness Hochreiter & Schmidhuber, 1997 defined a flat minimizer as a (large) region in weight space with the property that each weight vector from that region leads to similar small error, and propose an algorithm to search such a minimizer, called flat minimum search. Similarly, we can say that a minimizer is sharp when the loss function has a (very) high condition number near it, at any point in its basin of attraction. Keskar et al., 2017 characterize flatness using eigenvalues of the Hessian, and empirically show that  large-batch methods make the minimizers of training and testing functions sharp while small batch methods make them flat. Hochreiter & Schmidhuber, 1997 and Keskar et al., 2017 methods to quantify sharpness are not invariant to symmetries in the network, so, are not sufficient to determine generalization ability, as show by Dinh et al., 2017, Neyshabur et al., 2017.

Related phenomena Grokking can be seen as an extreme gradient starvation phenomenon [Pezeshki et al., 2021] or a very slow double descent [Nakkiran et al., 2020], for which the model only manages to learn the features necessary for generalization in the terminal phases of training [Pezeshki et al., 2021]. In fact, we observe that the
validation loss exhibits a double descent behaviour with an initial decrease, then a growth before rapidly decreasing to zero when the model groks. Grokking can also be seen as a phase transition (see [Nanda & Lieberum, 2022, Žunkovič & Ilievski, 2022]). Regarding catastrophic forgetting, there is a link between it when learning many tasks and the sharpness of the optimum for each task;

so that the slightest update to one task makes the optimum escape from its basin of attraction for the other tasks.

Mirzadeh et al., 2020, Mirzadeh et al., 2021 formalize this.

§ SUMMARY AND DISCUSSION

We made the following observations:

1. The memorization phase is characterized by a perturbed landscape, and it is separated from comprehension by a perturbed valley of bad solutions.
Small data results in the slow progression of SGD in this region, causing a delay in generalization. During the comprehension phase, the loss and accuracy of training and validation show a periodic perturbation. Thilak et al., 2022 named a related phenomenon the slingshot mechanism. We found that these perturbation points are characterized at the level of loss (resp. accuracy) by a sudden increase-decrease  (resp. decrease-increase), at the level of the model weights by a sudden variation of the relative cosine similarity, and at the level of the loss landscape by obstacles. This last point goes against what Goodfellow & Vinyals, 2015 observed, namely that a variety of state-of-the-art neural networks never encounter any significant obstacles from initialization to solution. The slingshot mechanism also contradicts the idea that SGD spends most of its time exploring the flat region at the bottom of the valley surrounding a flat minimizer [Goodfellow & Vinyals, 2015], since it goes with the model from confusion to the terminal phase of training, even after the model generalized, for many datasets.

2. The Hessian of the grokking loss function is characterized by larger condition numbers, leading to a slower convergence of gradient descent. We observed that more than 98% of the total variance in the parameter space occurs in the first 2 PCA modes much smaller than the total number of weights, suggesting that the optimization dynamics are embedded in a low-dimensional space [Li et al., 2018, Feng & Tu, 2021]. Moreover, the model remains in a lazy training regime [Chizat et al., 2019, Berner et al., 2021] most of the time, as the measure of cosine distance between the model weights from one training step to the next remains almost constant, except at the slingshot location.

From the point of view of the landscape, grokking seems a bit clearer: landscape geometry
has an effect on generalization, and can allow in the early stages of training to know if the model will generalize or not by  just looking at a microscopic quantity characteristic of that landscape like the empirical risk.

Limitations Some limitations  that would require further scrutiny include:

1. We view the loss surface with a dramatic reduction in dimensionality, and we need to be careful in interpreting these graphs.

2. We are only optimizing an unbiased estimate of the total loss function (added up across all training examples), whose structure may be different from the global loss function, so that the delay in generalization is an effect of each individual term in the loss function, or of the noise induced by the sampling of these terms, or of the stochasticity in the initialization of weights, generation of dropout masks, etc [Goodfellow & Vinyals, 2015]. More research needs to be done to find out if it is indeed the fact that the loss surface is perturbed that causes the generalization in the long run.

Perspectives The future direction is to study what a feature learning-based explanation of double descent [Pezeshki et al., 2021] has in common with grokking and to see if it is possible to design a specific learning curve as Chen et al., 2021 derived for the generalization curve. Also, this work focused on the detection of grokking however ultimately pave the way to also induce grokking. Another path of development is to move from arithmetic tasks with deterministic answers

to language models and computer vision with more probabilistic answers.

§ ACKNOWLEDGMENTS

We acknowledge the support from Canada CIFAR AI Chair Program and from the Canada Excellence Research Chairs Program.
This research used computing resources   on  Summit supercomputer at the Oak Ridge Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC05-00OR22725. The allocation was granted to  the INCITE project on “Scalable Foundation Models for Transferable Generalist AI”.

[Allen-Zhu et al., 2019]
Allen-Zhu, Z., Li, Y., and Liang, Y.
Learning and generalization in overparameterized neural networks,
going beyond two layers.
ArXiv, abs/1811.04918, 2019.

[Arora et al., 2018]
Arora, S., Ge, R., Neyshabur, B., and Zhang, Y.
Stronger generalization bounds for deep nets via a compression
CoRR, abs/1802.05296, 2018.
URL <http://arxiv.org/abs/1802.05296>.

[Belkin & Niyogi, 2001]
Belkin, M. and Niyogi, P.
Laplacian eigenmaps and spectral techniques for embedding and
In Dietterich, T., Becker, S., and Ghahramani, Z. (eds.),
Advances in Neural Information Processing Systems, volume 14. MIT
Press, 2001.

[Berner et al., 2021]
Berner, J., Grohs, P., Kutyniok, G., and Petersen, P.
The modern mathematics of deep learning.
arXiv preprint arXiv: Arxiv-2105.04026, 2021.

[Bubeck et al., 2015]
Bubeck, S. et al.
Convex optimization: Algorithms and complexity.
Foundations and Trends® in Machine Learning,
80 (3-4):0 231–357, 2015.

[Chen et al., 2021]
Chen, L., Min, Y., Belkin, M., and Karbasi, A.
Multiple descent: Design your own generalization curve.
Advances in Neural Information Processing Systems,
34:0 8898–8912, 2021.

[Chizat et al., 2019]
Chizat, L., Oyallon, E., and Bach, F.
On lazy training in differentiable programming.
Advances in Neural Information Processing Systems, 32, 2019.

[Cohen et al., 2021]
Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A. S.
Gradient descent on neural networks typically occurs at the edge of
International Conference On Learning Representations, 2021.

[David J.C. & Zoubin, 2005]
David J.C., M. and Zoubin, G.
Comments on ‘maximum likelihood estimation of intrinsic
dimension’ by elizaveta levina and peter bickel (2004).
URL <http://www.inference.org.uk/mackay/dimension/>.

[Dinh et al., 2017]
Dinh, L., Pascanu, R., Bengio, S., and Bengio, Y.
Sharp minima can generalize for deep nets.
In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th
International Conference on Machine Learning, ICML 2017, Sydney, NSW,
Australia, 6-11 August 2017, volume 70 of Proceedings of Machine
Learning Research, pp.  1019–1028. PMLR, 2017.
URL <http://proceedings.mlr.press/v70/dinh17b.html>.

[Donoho & Grimes, 2003]
Donoho, D. L. and Grimes, C.
Hessian eigenmaps: Locally linear embedding techniques for
high-dimensional data.
Proceedings of the National Academy of Sciences, 1000
(10):0 5591–5596, 2003.
URL <https://www.pnas.org/doi/abs/10.1073/pnas.1031596100>.

[Dziugaite & Roy, 2017]
Dziugaite, G. K. and Roy, D. M.
Computing nonvacuous generalization bounds for deep (stochastic)
neural networks with many more parameters than training data.
In Elidan, G., Kersting, K., and Ihler, A. T. (eds.),
Proceedings of the Thirty-Third Conference on Uncertainty in Artificial
Intelligence, UAI 2017, Sydney, Australia, August 11-15, 2017. AUAI
Press, 2017.
URL <http://auai.org/uai2017/proceedings/papers/173.pdf>.

[Facco et al., 2017]
Facco, E., d'Errico, M., Rodriguez, A., and Laio, A.
Estimating the intrinsic dimension of datasets by a minimal
neighborhood information.
Scientific Reports, 70 (1), sep 2017.
URL <https://doi.org/10.1038

[Feng & Tu, 2021]
Feng, Y. and Tu, Y.
The inverse variance–flatness relation in stochastic gradient
descent is critical for finding flat minima.
Proceedings of the National Academy of Sciences, 1180
(9):0 e2015617118, 2021.
URL <https://www.pnas.org/doi/abs/10.1073/pnas.2015617118>.

[Goodfellow & Vinyals, 2015]
Goodfellow, I. J. and Vinyals, O.
Qualitatively characterizing neural network optimization problems.
In Bengio, Y. and LeCun, Y. (eds.), 3rd International
Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May
7-9, 2015, Conference Track Proceedings, 2015.
URL <http://arxiv.org/abs/1412.6544>.

[Graves et al., 2013]
Graves, A., Mohamed, A., and Hinton, G. E.
Speech recognition with deep recurrent neural networks.
CoRR, abs/1303.5778, 2013.
URL <http://arxiv.org/abs/1303.5778>.

[He et al., 2015]
He, K., Zhang, X., Ren, S., and Sun, J.
Deep residual learning for image recognition.
CoRR, abs/1512.03385, 2015.
URL <http://arxiv.org/abs/1512.03385>.

[Herrmann et al., 2022]
Herrmann, L., Granz, M., and Landgraf, T.
Chaotic dynamics are intrinsic to neural network training with SGD.
In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.),
Advances in Neural Information Processing Systems, 2022.
URL <https://openreview.net/forum?id=ffy-h0GKZbK>.

[Hinton, 2012]
Hinton, G.
Neural networks for machine learning. coursera, video lectures, 2012.

[Hjorth, 1970]
Hjorth, B.
Eeg analysis based on time domain properties.
Electroencephalography and Clinical Neurophysiology,
290 (3):0 306–310, 1970.
ISSN 0013-4694.

[Hochreiter & Schmidhuber, 1997]
Hochreiter, S. and Schmidhuber, J.
Flat Minima.
Neural Computation, 90 (1):0 1–42, 01 1997.
ISSN 0899-7667.
URL <https://doi.org/10.1162/neco.1997.9.1.1>.

[Im et al., 2016]
Im, D. J., Tao, M., and Branson, K.
An empirical analysis of the optimization of deep network loss
arXiv preprint arXiv: Arxiv-1612.04010, 2016.

[Jastrzębski et al., 2018]
Jastrzębski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio, Y.,
and Storkey, A.
Three factors influencing minima in sgd.
ICLR, 2018.

[Jiang et al., 2019]
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S.
Fantastic generalization measures and where to find them.
International Conference On Learning Representations, 2019.

[Keskar et al., 2017]
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P.
On large-batch training for deep learning: Generalization gap and
sharp minima.
In 5th International Conference on Learning Representations,
ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track
Proceedings. OpenReview.net, 2017.
URL <https://openreview.net/forum?id=H1oyRlYgg>.

[Kingma & Ba, 2014]
Kingma, D. P. and Ba, J.
Adam: A method for stochastic optimization.
International Conference On Learning Representations, 2014.

[Kok & Sandrock, 2009]
Kok, S. and Sandrock, C.
Locating and characterizing the stationary points of the extended
rosenbrock function.
Evol. Comput., 170 (3):0 437–453, sep 2009.
ISSN 1063-6560.
URL <https://doi.org/10.1162/evco.2009.17.3.437>.

[Krizhevsky et al., 2012]
Krizhevsky, A., Sutskever, I., and Hinton, G. E.
Imagenet classification with deep convolutional neural networks.
In Pereira, F., Burges, C., Bottou, L., and Weinberger, K. (eds.),
Advances in Neural Information Processing Systems, volume 25. Curran
Associates, Inc., 2012.

[Lee et al., 2016]
Lee, J. D., Simchowitz, M., Jordan, M. I., and Recht, B.
Gradient descent converges to minimizers.
arXiv preprint arXiv: Arxiv-1602.04915, 2016.

[Levina & Bickel, 2004]
Levina, E. and Bickel, P.
Maximum likelihood estimation of intrinsic dimension.
In Saul, L., Weiss, Y., and Bottou, L. (eds.), Advances in
Neural Information Processing Systems, volume 17. MIT Press, 2004.

[Lewkowycz et al., 2020]
Lewkowycz, A., Bahri, Y., Dyer, E., Sohl-Dickstein, J., and Gur-Ari, G.
The large learning rate phase of deep learning: the catapult
arXiv preprint arXiv: Arxiv-2003.02218, 2020.

[Li et al., 2018]
Li, H., Xu, Z., Taylor, G., Studer, C., and Goldstein, T.
Visualizing the loss landscape of neural nets.
In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K.,
Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural
Information Processing Systems 31: Annual Conference on Neural Information
Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal,
Canada, pp.  6391–6401, 2018.

[Liu et al., 2022]
Liu, Z., Kitouni, O., Nolte, N. S., Michaud, E. J., Tegmark, M., and Williams,
Towards understanding grokking: An effective theory of representation
ArXiv, abs/2205.10343, 2022.

[Livni et al., 2014]
Livni, R., Shalev-Shwartz, S., and Shamir, O.
On the computational efficiency of training neural networks.
CoRR, abs/1410.1141, 2014.
URL <http://arxiv.org/abs/1410.1141>.

[Mirzadeh et al., 2020]
Mirzadeh, S. I., Farajtabar, M., Pascanu, R., and Ghasemzadeh, H.
Understanding the role of training regimes in continual learning.
Advances in Neural Information Processing Systems,
33:0 7308–7320, 2020.

[Mirzadeh et al., 2021]
Mirzadeh, S. I., Chaudhry, A., Hu, H., Pascanu, R., Gorur, D., and Farajtabar,
Wide neural networks forget less catastrophically.
ICML, 2021.

[Nakkiran et al., 2020]
Nakkiran, P., Kaplun, G., Bansal, Y., Yang, T., Barak, B., and Sutskever, I.
Deep double descent: Where bigger models and more data hurt.
In 8th International Conference on Learning Representations,
ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
URL <https://openreview.net/forum?id=B1g5sA4twr>.

[Nanda & Lieberum, 2022]
Nanda, N. and Lieberum, T.
A mechanistic interpretability analysis of grokking.
Alignment Forum, Aug 2022.

[Neyshabur et al., 2017]
Neyshabur, B., Bhojanapalli, S., McAllester, D., and Srebro, N.
Exploring generalization in deep learning.
In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus,
R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural
Information Processing Systems 30: Annual Conference on Neural Information
Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp.  5947–5956, 2017.

[Pezeshki et al., 2021]
Pezeshki, M., Kaba, O., Bengio, Y., Courville, A. C., Precup, D., and Lajoie,
Gradient starvation: A learning proclivity in neural networks.
In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan,
J. W. (eds.), Advances in Neural Information Processing Systems,
volume 34, pp.  1256–1272. Curran Associates, Inc., 2021a.

[Pezeshki et al., 2021]
Pezeshki, M., Mitra, A., Bengio, Y., and Lajoie, G.
Multi-scale feature learning dynamics: Insights for double descent.
icml, 2021b.

[Polyak, 1964]
Polyak, B.
Some methods of speeding up the convergence of iteration methods.
USSR Computational Mathematics and Mathematical Physics,
40 (5):0 1–17, 1964.
ISSN 0041-5553.

[Pope et al., 2021]
Pope, P., Zhu, C., Abdelkader, A., Goldblum, M., and Goldstein, T.
The intrinsic dimension of images and its impact on learning.
arXiv preprint arXiv:2104.08894, 2021.

[Power et al., 2022]
Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V.
Grokking: Generalization beyond overfitting on small algorithmic
arXiv preprint arXiv:2201.02177, 2022.

[Riedmiller & Braun, 1993]
Riedmiller, M. and Braun, H.
A direct adaptive method for faster backpropagation learning: the
rprop algorithm.
In IEEE International Conference on Neural Networks, pp.  586–591 vol.1, 1993.

[Roweis & Saul, 2000]
Roweis, S. T. and Saul, L. K.
Nonlinear dimensionality reduction by locally linear embedding.
Science, 2900 (5500):0 2323–2326, 2000.

[Shwartz-Ziv & Tishby, 2017]
Shwartz-Ziv, R. and Tishby, N.
Opening the black box of deep neural networks via information.
arXiv preprint arXiv: Arxiv-1703.00810, 2017.

[Silver et al., 2016]
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., van den Driessche,
G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M.,
Dieleman, S., Grewe, D., Nham, J., Kalchbrenner, N., Sutskever, I.,
Lillicrap, T., Leach, M., Kavukcuoglu, K., Graepel, T., and Hassabis, D.
Mastering the game of go with deep neural networks and tree search.
Nature, 529:0 484–503, 2016.

[Smith & Topin, 2017]
Smith, L. N. and Topin, N.
Exploring loss function topology with cyclical learning rates.
arXiv preprint arXiv: Arxiv-1702.04283, 2017.

[Tenenbaum et al., 2000]
Tenenbaum, J. B., de Silva, V., and Langford, J. C.
A global geometric framework for nonlinear dimensionality reduction.
Science, 2900 (5500):0 2319–2323, 2000.

[Thilak et al., 2022]
Thilak, V., Littwin, E., Zhai, S., Saremi, O., Paiss, R., and Susskind, J. M.
The slingshot mechanism: An empirical study of adaptive optimizers
and the {Grokking Phenomenon}.

[Vaswani et al., 2017]
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N.,
Kaiser, L., and Polosukhin, I.
Attention is all you need.
CoRR, abs/1706.03762, 2017.
URL <http://arxiv.org/abs/1706.03762>.

[Zhang et al., 2017]
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O.
Understanding deep learning requires rethinking generalization.
ArXiv, abs/1611.03530, 2017.

[Žunkovič & Ilievski, 2022]
Žunkovič, B. and Ilievski, E.
Grokking phase transitions in learning local rules with gradient
arXiv preprint arXiv: Arxiv-2210.15435, 2022.

§.§ Training phases

Previous work [Power et al., 2022, Liu et al., 2022] used the terms confusion, memorization and comprehension in the phase diagram based on different hyperparameters, but
in this paper they also refer to the phases along the training trajectory. Conventional deep neural networks training is split into two phases [Shwartz-Ziv & Tishby, 2017, Nakkiran et al., 2020, Feng & Tu, 2021]. A first phase during which the network learns a function with a small generalization gap and a second phase during which the network starts to overfit the data, leading to an increase in the test error. However, Nakkiran et al., 2020 show that in some regimes, the test error decreases again and can reach a lower value at the end of the training compared to the first minimum, suggesting potential training phases to exploit. Feng & Tu, 2021 distinguish two phases, an initial fast learning phase where the loss function decreases quickly and sometimes abruptly, followed by an exploration phase

when the training error reaches its minimum value and the overall loss still decreases, but much more slowly and gradually.  However, considering just two phases does not allow to study grokking properly, as the whole ingredient of grokking is in the memorization phase and in how the model gets in this phase and out.

[$r = 0.5$]
[$r = 0.6$]

On (a) and (b), the first figure (top) represents the validation accuracy (%) at the end of the training ($20k$ epochs), and the second figure (bottom) represents the epoch at which grokking happens ($-1$ is when we don't have grokking). On the x-axis we have the weight decay strength, and on the y-axis we have the learning rate.

§.§ Non-grokking

$t_4(r)$ is a transition step from memorization to generalization for each training data fraction $r >= 0.2$. The blue curve shows the empirically obtained values of $t_4(r)$ (y-axis) as a function of $r$ (x-axis), and the orange curve represents an estimate of the form $t_4(r) = a r ^{-\gamma} + b$ fitted on the obtained data, with $(\gamma, a, b) = (7.73, 1.09 \times 10^{15}, 1442.63)$ in this

Left) The first figure (top) represents the validation accuracy (%) at the end of the training ($10k$ steps), and the second figure (bottom) represents the spectral energy (activity) in the training loss for the first 400 training steps ($r = 0.5$). On the x-axis we have the weight decay strength, and on the y-axis we have the learning rate. Right) On the x-axis we have the (normalize) activity (400 steps), and on the y-axis we have the validation accuracy (%), for different value of $r$.
On the x-axis we have the weight decay strength, and on the y-axis we have the learning rate. A similarity is observed between the oscillation patterns in the training loss during the initial stages of training and the validation accuracy. This suggests that the spectral signature can be used as an indicator or proxy for the upcoming grokking phenomenon.
The highest degree of generalization is typically observed when using small learning rates and small weight decay. While large learning rates may increase oscillations, this does not directly lead to grokking and is not necessarily evident in the early stages of training. Instead, such effects become more noticeable near the basin of attraction of the minimum.

For (a)-(d), the first figure (top) represents the validation
accuracy ($\%$) at the end of the training (10k steps), and
the second figure (bottom) represents the spectral energy
(activity) in the training loss for the first 400 training steps. On the x-axis we have the weight decay strength,
and on the y-axis we have the learning rate. A similarity is observed
between the oscillation patterns in the training loss during
the initial stages of training and the validation accuracy.

Although it is easy to identify grokking, it is very difficult to give a formal definition of its opposite since nothing contradicts the fact that if the model is allowed to train indefinitely it will not eventually generalize. Knowing the hyperparameters (learning rate, weight decays strength, ...) that allow to have grokking (

let's note by $H_g$ the set of values of such hyperparameters), we use them to train our model for different data sizes. Then we fit a function that predicts $t_4$, the generalization step, for each training data fraction $r$.  We then used the expression of $t_4(r)$ estimated as a proxy of the number of training steps needed to get the generalization for each $r$. That is, if for other hyper-parameters that are not in $H_g$, and for a given training data fraction $r$, we train a model for $t_4(r)+\epsilon$ (we used $\epsilon = 1k$) steps and have no generalization, we can stop the training. This does not necessarily imply that there will be no grokking if the model is left in training for longer.  The limit of this approach is that $t_4(r)$ is just an empirical law estimated for $r \ge r_{min} > 0$, which collapses when $r \rightarrow 0$.

§.§ The over-parameterization ratio and the effect of the optimizer parameters

Evolution of the validation accuracy for modular addition with $r \in \{20, 25, \dots, 90\} \%$. More training data leads to faster grokking.

In general, more data leads to faster grokking (figure

), i.e. $t_2(r)$ and $t_4(r)$ are a decreasing functions of $r$. Empirically, $t_4$ follows a power law of the form $t_4(r) = a r ^{-\gamma} + b$. This was first predicted by Žunkovič & Ilievski, 2022. For $r \ge r_{min}$, we found $(\gamma, a, b) = (7.73, 1.09 \times 10^{15}, 1442.63)$ for modular addition (figure <ref>.b), and $(\gamma, a, b) = (1.18, 1.85 \times 10^{6}, 0.0)$ for multiplication in $S_5$, in $H_g$.

Smaller learning rates require more training steps for convergence whereas larger learning rates result in rapid changes and require fewer training epochs. But a learning rate that is too large can cause the model to converge too quickly to a suboptimal solution, whereas a learning rate that is too small can cause the process to get stuck. This is well verified with grokking, because the more we increase the learning step, the faster we observe grokking, up to a threshold that depends on the value of the weight decay strength used.

§.§ Spectral signature

Figures and <ref> and <ref> show a similarity between the oscillation in the training loss in the early phases of training and the validation accuracy for many values of $r$.

§ LOSS LANDSCAPE

For $r = 0.3$, figure <ref> shows the 1D projection of the grokking loss surface for a single epoch of training (just after the grokking step), while figures <ref> ($\vec{\delta}_t \propto \theta^* - \theta_t$), <ref> ($\vec{\delta}_t \propto \theta_{t+1} - \theta_t$), <ref> ($\vec{\delta}_t \propto \theta_0 - \theta_t$) and <ref> (random $\vec{\delta}_t$) show it for different training epochs.

[$r=30\%$ ($S_5$) ]
[$r=30\%$ ($+$)
[$r=60\%$ ($+$)

[$r=80\%$ ($+$)
1D projection of the grokking loss surface.
The x-axis is $\alpha \in [-3, 3]$, and the y-axis are the loss (left, blue) and accuracy (right, red) at $\theta^* + \alpha \frac{\theta_0 - \theta^*}{\| \theta_0 - \theta^*\|}$ (with

filter-wise normalization [Li et al., 2018]), where $\theta_0$ is the initial parameter and $\theta^*$ the minimizer found by the optimization algorithm.

(a) is obtained with multiplication in $S_5$ as operator, and the others with modular addition. $r$ is the percentage of total data used for training. We can see that the 1-D subspace from initial to final parameters and from one minimizer to another contains many difficult and exotic structures. In (a), we have two minimizers of the training loss (solid line, -), but only one of them also minimizes the validation loss (dashed line, - -).

in 1, 99, 350, 1186, 1970

Random $\vec{\delta}$ direction ($r = 0.8$). In this direction, no obstacles are visible along the trajectory, and the surface is rather well-conditioned.

[Loss surface : $f_t(\alpha) = Loss(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

[Accuracy surface : $f_t(\alpha) = Acc(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

1D projection of the grokking loss/accuracy surface in a random direction $\vec{\delta}$. In this direction, no obstacles are visible along the trajectory, and the surface is rather well-conditioned.

in 0, 2, 100, 200, 350, 975, 1187, 1387, 1587, 1968

From the initialization point to any step $t$, $\vec{\delta}_t = \theta_0 - \theta_t$ ($r = 0.8$).
In this direction, each point considered behaves locally as a minimum, and as one approaches the optimum, local minima becomes preponderant. Although the model generalizes at  $t<1k$, the shape of the landscape continues to evolve, and towards the end, we observe two local minima near the optimum: one on the left, very flat, but separated from the optimum by a mountainous spike, and a second one on the right, separated by a more easily escaped mountain.

[Loss surface : $f_t(\alpha) = Loss(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

[Accuracy surface : $f_t(\alpha) = Acc(\theta_t + \alpha \vec{\delta}_t)$ for each epoch $t$]

1D projection of the grokking loss/accuracy surface from the initialization point to any step $t$, $\vec{\delta}_t \propto \theta_0 - \theta_t$.
In this direction, each point considered behaves locally as a minimum, and as one approaches the optimum, local minima become preponderant.

§.§ Loss surface curvature

The approach described in the previous section allows us to visualize the loss surface under a dramatic dimensionality reduction, and we need to be careful how we interpret the plots. One way to measure the level of convexity in a loss function is to compute the principal curvatures, which are simply eigenvalues of the Hessian. Since we are in a very high dimension, it would be impractical to compute all the eigenvalues of the Hessian, or simply to compute the Hessian itself. To avoid this dimensional problem, we estimated the curvature through the condition number of the Hessian of $L(\theta)$ at each step $t$, which is simply the ratio $\lambda_{min}(t) / \lambda_{max}(t)$, where $\lambda_{min}(t)$ and $\lambda_{max}(t)$ are, respectively, the minimum and maximum eigenvalues of the Hessian of $L_t$.

The condition number is a direct measure of pathological curvature. Larger condition numbers imply slower convergence of gradient descent. Li et al., 2018 compute the above ratio
at each point of the loss surface and observe that the convex-looking regions in the surface plots correspond to regions with insignificant negative eigenvalues, while chaotic regions contain large negative curvatures; and that for convex-looking surfaces, the negative eigenvalues remain extremely small.

To measure the level of curvature of the loss function, we compute the maximum (${\lambda}_{max}$) and minimum (${\lambda}_{min}$) eigenvalue of its Hessian (figure <ref>). We observe that there is no significant negative curvature in the trajectory. The curvature remains generally positive and is greatly disturbed at the slingshot points (${\lambda}_{min}$ remains in general close to 0 while ${\lambda}_{max}$ is large, but during slingshot ${\lambda}_{min}$ becomes negative). Our loss is ${\lambda}_{max}$-smooth, i.e. its gradient is ${\lambda}_{max}$-Lipschitz. It is known from optimization literature that a function with bounded Hessian eigenvalues has a gradient that tends to decay when the parameter gets closer to the minimum, in contrast to a non-smooth one that generally has abrupt bends at the minimum, which causes significant oscillations for gradient descent [Bubeck et al., 2015].

Let $\{\lambda_t(i)\}_i$ be the eigenvalues of $\mathcal{H}_t$ and $\{v_t(i)\}_i$ the associated eigenvectors. For a very small step size $\epsilon_t$ of SGD, $L_{t+1} - L_t \approx - \epsilon_t \| G_t \|^2 +  \frac{1}{2} \epsilon_t^2 G_t^T \mathcal{H}_t G_t - o(\epsilon_t^2 \|G_t\|^2)$. Futher, if $\lambda_t(i)  > 2/\epsilon_t$, we get $- \epsilon_t \| G_t \|^2 +  \frac{1}{2} \epsilon_t^2 G_t^T \mathcal{H}_t G_t > 0$.

In fact, Let $\{\lambda_t(i)\}_i$ be the eigenvalues of $\mathcal{H}_t$ and $\{v_t(i)\}_i$ the associated eigenvectors. For a very small step size $\epsilon_t$ of SGD, $L_{t+1} - L_t \approx - \epsilon_t \| G_t \|^2 +  \frac{1}{2} \epsilon_t^2 G_t^T \mathcal{H}_t G_t - o(\epsilon_t^2 \|G_t\|^2)$. Futher, if $\lambda_t(i)  > 2/\epsilon_t$, we get $- \epsilon_t \| G_t \|^2 +  \frac{1}{2} \epsilon_t^2 G_t^T \mathcal{H}_t G_t > 0$ [This comes come from the fact that near $\theta_t$, $L(\theta) \approx L_t + (\theta - \theta_t)^T G_t +  \frac{1}{2} (\theta - \theta_t)^T \mathcal{H}_t (\theta - \theta_t) + o(\|\theta - \theta_t\|^2)$. Taking $\theta = \theta_t - \epsilon_t G_t$ gives the first approximation. The last inequality is obtained with $G_t^T \mathcal{H}_t G_t = \sum_{i} \lambda_t(i)
\langle G_t, v_t(i) \rangle^2$ and $\sum_{i} \langle G_t, v_t(i) \rangle ^2  = \| G_t \|^2$.

When $\mathcal{H}_t$ has some large positive eigenvalues (i.e., high-curvature directions) and some eigenvalues close to 0 (i.e., low-curvature directions), gradient descent bounces back and forth in high-curvature directions and makes slow progress in low-curvature directions. In this case, the optimization problem has an ill-conditioned curvature. Furthermore, if the loss function near $\theta_t$ has a high condition number, that is, very small steps cause an increase in the cost function (for example, if $\theta_t$ is a very sharp minimum surrounded by high loss regions),

the optimization problem becomes also ill-conditioned. During training,

if the gradient norm does not shrink but $G_t^T \mathcal{H}_t G_t$ increases in order of magnitude, learning can become very slow despite a strong gradient.

The above observation ($\lambda_t(i) > 2\epsilon_t^{-1}$)
is similar to what Herrmann et al., 2022 [Theorem 2.1] defines for the eigenvalues of a positive-definite $\mathcal{H}_t$, as a locally chaotic training behaviour of the Local Lyapunov Exponents. They show evidence that neural network training is intrinsically locally chaotic due to the negative eigenspectrum of the Hessian, and that network training with SGD exhibits globally edge-chaotic behaviour. This observation is also linked to the progressive sharpening phenomenon [Cohen et al., 2021] in which $max_{i} \lambda_t(i)$ increases and reaches a value that is equal to or slightly larger than $2\epsilon_t^{-1}$, leading the model to enter an Edge of Stability regime where loss shows non-monotonic training behaviour over short time spans [Thilak et al., 2022].

Condition number of the Hessian over training epoch for modular addition (top, $r=0.3$, bottom, $r=0.8$).

We can see that there is no significant negative curvature in the trajectory (figure <ref>). The curvature remains in general positive and is greatly disturbed at the slingshot points.

We can see that there is no significant negative curvature in the trajectory (figure <ref>). The curvature remains in general positive, and is greatly disturbed at the slingshot points. Our loss is ${\lambda}_{max}$-smooth, i.e. its gradient is ${\lambda}_{max}$-Lipschitz. It is known from optimization literature that a function with bounded Hessian eigen-values has the gradient that tend to decay when the parameter gets closer to the minimum, in constrast to non-smooth one that generaly have abrupt bends at the minimum, which cause significant oscillations for gradient descent [Bubeck et al., 2015].

Let $\{\lambda_t(i)\}_i$ be the eigenvalues of $\mathcal{H}_t$ and $\{v_t(i)\}_i$ the associated eigenvectors. For a very small step size $\epsilon_t$ of SGD, $L_{t+1} - L_t \approx - \epsilon_t \| G_t \|^2 +  \frac{1}{2} \epsilon_t^2 G_t^T \mathcal{H}_t G_t - o(\epsilon_t^2 \|G_t\|^2)$. Futher, if $\lambda_t(i)  > 2/\epsilon_t$, we get $- \epsilon_t \| G_t \|^2 +  \frac{1}{2} \epsilon_t^2 G_t^T \mathcal{H}_t G_t > 0$.

In fact, near $\theta_t$, $L(\theta) \approx L_t + (\theta - \theta_t)^T G_t +  \frac{1}{2} (\theta - \theta_t)^T \mathcal{H}_t (\theta - \theta_t) + o(\|\theta - \theta_t\|^2)$. Taking $\theta = \theta_t - \epsilon_t G_t$ gives the first approximation. The last inequality is obtained with $G_t^T \mathcal{H}_t G_t = \sum_{i} \lambda_t(i)
\langle G_t, v_t(i) \rangle^2$ and $\sum_{i} \langle G_t, v_t(i) \rangle ^2  = \| G_t \|^2$.

When $\mathcal{H}_t$ has some large positive eigenvalues (i.e., high-curvature directions) and some eigenvalues close to 0 (i.e., low-curvature directions), gradient descent bounces back and forth in high-curvature directions and makes slow progress in low-curvature directions. In this case, the optimization problem has an ill-conditioned curvature. Furthermore, if the loss function near $\theta_t$ has a high condition number, that is, very small steps cause an increase in the cost function (for example, if $\theta_t$ is a very sharp minimum surrounded by high loss regions),

the optimization problem becomes also ill-conditioned. During training,

the gradient norm does not shrink but $G_t^T \mathcal{H}_t G_t$ increases in order of magnitude, learning can become very slow despite a strong gradient.

The last observation of proposition <ref> ($\lambda_t(i) > 2/\epsilon_t$)
is similar to what Herrmann et al., 2022 [Theorem 2.1] defines for the eigenvalues of a positive-definite $\mathcal{H}_t$, as a locally chaotic training behavior of the Local Lyapunov Exponents. They show evidence that neural network training is intrinsically locally chaotic due to the negative eigenspectrum of the Hessian, and that network training with SGD exhibits globally edge-chaotic behavior. This observation is also link to the progressive sharpening phenomenon [Cohen et al., 2021] in which $max_{i} \lambda_t(i)$ increases and reaches a value that is equal to or slightly larger than $2\epsilon_t^{-1}$, leading the model to enter an Edge of Stability regime where loss shows non-monotonic training behavior over short time spans [Thilak et al., 2022].

§.§ Flatness

It is necessary to quantify the sharpness or the flatness of the landscape beyond the visual representations. Indeed, the flatness can contain nonlocal properties of landscape as it depends on the loss function in a finite neighborhood of a given point in parameter space, whereas the Hessian matrix only describes the local curvature of the landscape.
Let consider again the matrix $M = [\theta_t - \theta_T]_{1 \le t \le T - 1} \in \mathbb{R}^{(T-1) \times d} $ where $T$ is the total number of training steps. We applied PCA to the matrix $M$, select the $n$ most explanatory directions. For each PCA direction $i$ near a given point $\theta$ in the parameter space, we define a flatness $F_i(\theta)$ as the width of the region within which $L_i \le \beta \times L(\theta)$ where $\beta \ge 1$ is a hyperparameter. This is the method used by Feng & Tu, 2021 with $\beta$ equals to the Euler number. They find the two closest points $\theta^l_i < 0$ (left) and $\theta^r_i > 0$ (right) on each side of $\theta$, in the direction $i$, that satisfy $L(\theta^l) = L(\theta^r) \approx \beta \times L(\theta)$, and define the flatness parameter as $F_i(\theta) = \theta^r_i - \theta^l_i$. For a PCA direction $\vec{\delta}_i$, we can redefine this in our framework as

F_i(\theta) =  min\{\alpha > 0 \ | L(\theta + \alpha \vec{\delta}_i) \ge \beta \times L(\theta) \} +
min\{\alpha > 0 \ | L(\theta - \alpha \vec{\delta}_i) \ge \beta \times L(\theta) \}

The larger $F_i(\theta)$ is, the flatter is the landscape in the $i^{th}$ PCA direction.

Alternatives : use the trace of the Hessian
at the minima $Tr(H) = \sum_i \lambda_i$ (a lower value of $Tr(H)$ the wider the minima) or its Frobenius norm (higher values imply a sharper minimum) [Jastrzębski et al., 2018]. The largest eigenvalue and the Frobenius norm (approximated using finite difference method) of the Hessian can be used in place of the trace of the Hessian because calculating multiple eigenvalues to directly approximate the trace is too
computationally intensive.

TODO : plots and interpretation

§.§ Other landscape measures

TODO :
To see how far optimization strays from the primary linear subspace, we plot the norm of the residual of the parameter value after projecting the parameters at each epoch into the 1-D subspace using one of the following two approaches.

* Let $M = [\theta_t - \theta_T]_{1 \le t \le T - 1}$ where $T$ is the total number of training epochs. We applied PCA to the matrix $M$, select the $2$ most explanatory directions, then project each parameter $\theta_t$ on this these two directions to have $\alpha(t)$ and $\beta(t)$.

* Let $\vec{\delta}$ be a unit vector pointing in the direction $\theta^* - \theta_0$, and $\alpha(t) = (\theta_t - \theta_0)^T \vec{\delta}$, where $\theta_0$ is the initial parameter and $\theta^*$ the optimum find by our optimization procedure. It is this quantity that we try to quantify when we interpolate the loss from the initialization point to the final training point, in 1-D. Let $\vec{\eta}(t)$ be a second unit vector pointed in the direction $\theta_t - (\theta_0 + \alpha(t) \vec{\delta})$, and $\beta(t) = (\theta_t - \theta_0 - \alpha(t) \vec{\delta})^T \vec{\eta}(t)$ the norm of the residual of the projection of $\theta_t$ onto the line spanning initialization and solution [Goodfellow & Vinyals, 2015].

We will visualize how $\alpha(t)$ and $\beta(t)$ evolve jointly during training.

To see how far optimization strays from the primary linear subspace, we plot the norm of the residual of the parameter value after projecting the parameters at each epoch into the 1-D subspace using the following two approaches. Let $M = [\theta_t - \theta_T]_{1 \le t \le T - 1} \in \mathbb{R}^{(T-1) \times d}$ where $T$ is the total number of training steps. We applied PCA to the matrix $M$, select the $2$ most explanatory directions, then project each parameter $\theta_t$ on this these two directions to have $\alpha(t)$ and $\beta(t)$.

On figure <ref>, we plot the projection along the 2 first PCA axis from initialization to solution, $\alpha(t)$ and $\beta(t)$. More than 98% of the total variance in the parameter space occurs in the first 2 PCA modes much smaller than the total number of weights, suggesting that the optimization dynamics are embedded in a low-dimensional space [Li et al., 2018, Feng & Tu, 2021].

30, 35, 40, 45, 50, 55, 60, 65, 70, 75, 80,

[$ r = \tdp \%$]

Projection along the axis from initialization to solution, $\alpha(t)$ and $\beta(t)$. We can see that more than 98% of the total variance in the parameter space occurs in the first 2 PCA modes much smaller than the total number of weights, suggesting that the optimization dynamics are embedded in a low-dimensional space.

Gradient Norm $\| G_t \|^2$. Most of the time there is no significant gradient, which results in a weak progression of gradient descent.

Relative consine similarity ($r=0.3$ and $0.8$). Left) $cos(\theta_t$, $\theta_{t+1})$ Right) $cos(\theta_t$, $\theta_{0})$. Most of the time the weights progress very little, jumping when the gradient explodes, with a slingshot effect.

Also, the cosine similarity measure between the model weights from one epoch to the next remains almost constant, except at the slingshot location (figure <ref>). This allows us to see that our model passes through the anomalies observed above. It has been shown that with high probability over the initialization; the iterates of the gradient descent algorithm even stay in a small fixed neighbourhood of the initialization during training. Because the parameters only move very little, this type of training has also been coined lazy training <cit.>. These steps where the distances change abruptly are caused by spikes in the training gradient (figure <ref>).

§.§ The Slingshot mechanism

Above, we derive $\dot{L}(t) = - \| G(t) \|^2$ for a small enough step size. This relaxation property of gradient descent reflects the fact that the loss function cannot increase. However, we lose this property with certain accelerated or adaptative methods like Adam, or on ill-conditioned problems as we saw above with $L_{t+1} > L_t$ for $\epsilon_t \lambda_t(i) > 2$. In the case of grokking, the loss, even when it becomes zero, presents a sudden growth, followed by decay, and this is in a periodic way.
Thilak et al., 2022 have recently shown that this phenomenon is general to the optimization of deep neural networks. Between two slingshots, the gradient is almost zero, all eigenvalues of the Hessian are nonnegative, and there is one direction of very large curvature that dominates the others and the weights update, resulting in a small change in weight from one training step to another. The model thus seems to traverse a flat valley.

To reduce this phenomenon (i.e. the amplitude of the spikes), we clipped the gradient norm during training (using a threshold $\eta > 0$).  This had the effect of slowing down generalization, but not preventing it.

\dot{\theta} = \left\{
\begin{array}{ll}
- \frac{\eta}{\| G(t) \|}     G(t) & \mbox{if } \| G(t) \| \ge \eta \\
- G(t) & \mbox{otherwise.}
\end{array}
\right. (\eta > 0)

\Longrightarrow \dot{L} = \left\{
\begin{array}{ll}
- \eta \| G(t) \| & \mbox{if } \| G(t) \| \ge \eta \\
- \| G(t) \|^2 & \mbox{otherwise.}
\end{array}
\right. (\eta > 0)

The spikes are reduced as $\eta \rightarrow 0$, but remain visible as $\theta \rightarrow \theta^*$.  By reducing the learning rate extremely, the spikes also become less visible, but we also pay the same cost because the model takes more steps to generalize.

Effect of gradient clipping on slingshot

TODO : (1) figure (2) how this affect the signature

§.§ Three-dimensional visualizations

We can accomplish this by viewing a heatmap of the cost function in 2D (figure <ref>).

3D ($r=0.5$)

90, 85, 80, 75, 70, 65, 60, 55, 50,

40, 35,

1D projection of the grokking loss (top) and accuracy (bottom) surface for different value of the training data fraction $r$, for modular addition. The direction used for each training epoch $t$ is $\vec{\delta}_t = \theta^* - \theta_t$.

90, 85, 80, 75, 70, 65, 60, 55, 50,
40, 35

1D projection of the grokking loss (top) and accuracy (bottom) surface for different value of the training data fraction $r$, for multiplication in $S_5$. The direction used for each training epoch $t$ is $\vec{\delta}_t = \theta^* - \theta_t$.

30, 35, 40, 45, 50, 55, 60,  65, 70, 75,

in 1, 2, 3, 4, 5, 6, 7

From the initialization point (modular addition) : $\vec{\delta}_t = \theta_0 - \theta_t$

35, 40, 45, 50, 55, 60,  65, 70, 75,

85, 90

Condition number of the Hessian over training epoch for modular addition.

35, 40, 45, 50, 55, 60,  65, 70, 75,

85, 90

Gradient Norm for modular addition.

35, 40, 45, 50, 55, 60,  65, 70, 75,

85, 90

Relative cosine distance for modular addition. Left) $cos(\theta_t$, $\theta_{t+1})$ Right) $cos(\theta_t$, $\theta_{0})$