We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version.

Loss curves without weight decay?

For readability, we omitted the loss curves without weight decay from the figure. Overall, the final solution without weight decay closely matches the (regularized) state at time $t_2$​, and the corresponding loss curves align with those shown up to $t_2$​ before remaining essentially flat. We would be happy to include these curves in the revised version if the reviewer believes it would improve the presentation.

Clarify dependence on initialization scale

Using the additional page allowed for the camera-ready version, we will include a more detailed discussion on the influence of initialization scale. The initialization scale affects the behavior of the unregularized flow and consequently impacts the first phase of training under regularization. This scale plays a crucial role in the training dynamics, as extensively studied in the implicit bias literature. Although a complete understanding remains elusive, the common consensus—supported by many studies—is that small initialization scales correspond to the rich regime, where implicit bias guides the model toward interpolating solutions with smaller weight norms, which generally leads to better generalization. In contrast, large initialization scales lead to the NTK or lazy regime, where features remain nearly fixed during training. This regime behaves similarly to random feature models and typically results in interpolators with poorer generalization performance.

In our setting, this means that when starting from a small initialization scale, the point $\Phi( w(0) )$ reached after the first phase already has a small norm--potentially corresponding to a KKT point of Equation (7)--and thus no second “grokking” phase occurs (i.e., no significant change happens afterward). Conversely, with a large initialization scale, $\Phi(w(0))$ has a large norm, leading to substantial movement during the second grokking phase.

We will include additional experiments in the appendix of the revised version to illustrate the impact of initialization scale on the grokking phenomenon.

We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version. These questions/remarks are very high level and we share most of these concerns/questionings. We give below our “point of view” on these matters, more than fully supported answers.

Why does grokking matter?

In that direction, a recent position paper at ICML 2025 [1] argued that grokking may not be a particularly central phenomenon and suggested that the research community should be cautious about devoting excessive attention to it. However, even within that paper, the authors emphasize that studying such settings can contribute to a broader and deeper understanding of training dynamics in deep learning. We fully agree with this perspective and believe our work advances this broader goal. Our theoretical framework is general enough to capture meaningful and widely applicable behaviors in neural network training. While our main results are motivated by the grokking phenomenon, they also offer insights beyond it. Specifically:

1. Our analysis provides a general characterization of training dynamics in the small-regularization regime, which can be useful for understanding behavior in a variety of settings—not limited to grokking.
2. Although grokking is often discussed in the context of weight decay, our approach can easily be extended to other forms of regularization. In particular, we believe that some empirical observations reported in “The effect of SAM for deep networks at different stages of training” [2] for training with Sharpness-Aware Minimization (SAM), may also be interpreted through the lens of grokking, albeit driven by SAM-style regularization rather than standard weight decay.

This doesn’t seem conceptually novel

We agree that the intuition that “grokking is due to a first interpolating phase, followed by a second phase where regularization kicks in” has been suggested before, for instance in the OmniGrok paper. Since this idea is natural and makes sense, we were actually surprised not to find any formal mathematical treatment of this intuition in the literature. Our work aims to fill this gap. Moreover, while it is nice to say that “regularization kicks in,” as mathematicians, we appreciate having a more precise description. This is what we provide here, showing that this second phase exactly corresponds to a Riemannian gradient flow for the $\ell_2$​-norm on the interpolation manifold.

Additionally, we felt that this interpretation of grokking was not universally understood, and that the role of weight decay remained broadly unclear, as raised in our “role of weight decay” paragraph.

we've formalized the easy part, but the deeper mystery, except in special cases, is still open.

Of course, the question of generalization in deep learning is a longstanding and largely unresolved problem that has been extensively studied over many years. In this work, we build upon the commonly held belief that smaller norm weights tend to correlate with better generalization, while acknowledging that this relationship is far from definitive in general settings. Our goal is not to solve this fundamental question, but rather to leverage existing insights on generalization to provide an interpretation of grokking that complements our optimization analysis.

References

[1] Jeffares, A., & van der Schaar, M. Position: Not All Explanations for Deep Learning Phenomena Are Equally Valuable. In Forty-second International Conference on Machine Learning Position Paper Track.

[2] Andriushchenko, Maksym, and Nicolas Flammarion. "Towards understanding sharpness-aware minimization." International conference on machine learning. PMLR, 2022.

We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version.

1. How do you explain the sudden drop in validation loss observed in grokking?

We would like to clarify that the observed drop is not actually sudden, even in practice. It may appear so because, as in most grokking papers, results are plotted in log-time scale. In such plots, the "drop" phase occupies a much smaller portion of the x-axis than the preceding plateau, which can create the impression of a sudden change. However, when viewed on a linear time scale, the drop is more gradual. This is not unique to our experiments but is common across many grokking studies (e.g., Power et al., 2022; Liu et al., 2022c).

Furthermore, Proposition 2 in our paper gives a theoretical explanation for this. It predicts that the second phase of learning occurs within a time interval $[\varepsilon/\lambda, M/\lambda]$, where $\varepsilon$ is small and $M$ is large, both independent of $\lambda$. As a result, on a log scale, the drop appears within an interval of length $\ln⁡(M/\varepsilon)$, while the plateau spans $[\Theta(1),\Theta(1)/\lambda]$, corresponding to a length of $\ln(1/\lambda)$ on the same scale. This explains why the drop looks sudden relative to the plateau in log-scale plots. However, in linear time, the drop spans a length of a similar order of magnitude than the preceding plateau, giving a different visual impression.

2. Relation to Li et al (2021)

We acknowledge that Li et al. (2021) revived the “two-stage” gradient flow analysis, which is the foundation of our proof of Proposition 2. We do not claim novelty of this analytical framework. In fact, this approach dates back to Katzenberger (1990) and was further developed in more modern contexts by Fatkullin et al. (2010). The limit flow $\Phi$ is a standard object in continuous-time optimization and already appears in Katzenberger’s work. We will clarify these connections more explicitly in the revision. In particular, we will emphasize how Li et al. (2021) recently highlighted the power of this technique.

At a high level, our contribution differs from Li et al. (2021) in that we explain a slow drift induced by a deterministic regularization, whereas their analysis focuses on stochastic effects. Additionally, we relate this drift directly to the grokking phenomenon, which to our knowledge has not been previously addressed in this context. From a technical standpoint, our setting is simpler and allows for a more detailed analysis:

1. Rather than invoking Katzenberger (1990) as a black-box result, we provide a self-contained and accessible proof that builds on Falconer (1983), whose ideas also underlie the works of Katzenberger and Li et al.
2. Moreover, while prior analyses typically assume initialization close to the interpolation manifold, our approach handles general initialization. We rigorously analyze the first phase, where the flow approaches the manifold, and then establish the transition to the second phase (drift along the manifold). Dealing with this transition is non-trivial and, as we mention following Proposition 2, is the main technical novelty in this proof.

3. Experiments on real-world datasets.

Many empirical studies have already demonstrated the grokking phenomenon on real-world datasets. We reference several of these in the paragraph “Grokking in experimental works” in Section 2, and we do not believe that reproducing similar experiments would yield additional insights.

Our objective is not to contribute further real-world evidence, but rather to validate the coherence between our theoretical results and controlled toy simulations. These simplified settings allow us to precisely monitor the relevant quantities and illustrate our theoretical findings.

Naturally, we would be happy to include any additional references of interest if you have suggestions.

4. The theory relies on the $\lambda \to 0$ limit.

Our work indeed focuses on the $\lambda \to 0$ limit, as it allows for a simple and tractable analysis. Providing clean results with a constant $\lambda > 0$ would require stronger assumptions and likely lead to much more complex proofs.

Regarding the statement that

even theoretically, if $\lambda$ is a nonzero constant, the gradient flow in the first stage may not converge to the manifold at all,

although this is true, we can still show that the flow reaches a point that is close to the manifold, and upper bound the distance as a function of $\lambda$.

In particular, Equation (9) in the Appendix (proof of Proposition 1) provides an upper bound on the deviation between the regularized and unregularized flows. This allows us to bound the distance of the unregularized flow to the manifold for finite $\lambda$ and a carefully chosen time $t$ (as detailed in the proof of Lemma 2). While we acknowledge that this bound may not be tight enough for practical use, it does offer some theoretical control in the finite $\lambda$ setting.

Also note that in all our synthetic experiments, which are run for non-zero but small values of $\lambda$, the training loss goes down to nearly zero, indicating that the iterates nearly converge towards the manifold.

Besides, although our results are asymptotic, they appear to match what we observe empirically for fixed, small $\lambda$.

We hope our response clarifies the points raised and appreciate the reviewer’s consideration of these explanations.

We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version.

When do Assumptions 1, 2 hold?

In short: our assumptions are quite generic. They hold for most basic neural network architectures with smooth activations, up to some important caveats (in general, they might only hold locally, and after excluding certain degenerate configurations).

Assumption 1 is pretty mild: it holds for all architectures that use differentiable activations (i.e. that do not use ReLU). Although ReLU is a very common activation, many architectures now use differentiable variants of ReLU such as GeLU and thus satisfy Assumption 1. The ‘restrictive’ part of Assumption 1 is that it assumes that $w_{GF}$ is bounded: this does not hold when training with cross entropy loss (or other exponential tail classification losses) on separable datasets. Apart from this case (i.e. in regression tasks or when the convergence point does not interpolate the data), Assumption 1 holds.

Assumption 2 (Morse-Bott property) is the more restrictive one. It requires the Hessian matrix $\nabla^2 F(w)$ to have a constant rank on $\mathcal{M}$. In general, this is not true globally, as $\mathcal{M}$ can have degenerate points. However, in most cases it holds locally around almost every point of $\mathcal{M}$. For example:

• for the "diagonal neural network" toy model, it holds if we restrict to points with nonzero coordinates,
• for matrix completion, it does if we exclude matrices with degenerate rank,
• it also holds locally almost everywhere for wide MLPs, ConvNets and ResNets with smooth activations; see for instance Liu, Zhu and Belkin (2021)

In our work, we preferred to assume that the Morse-Bott property holds globally to derive simpler results and analysis. Similar global assumptions were considered by the works of Li et al (2021), Shalova et al (2024), Fatkullin et al (2010). Although it is not exactly true in practice, we believe it still allows to derive meaningful analyses. Note that we could extend our results to the local case, provided that the iterates remain in the non-degenerate region.

We will add more discussion about these assumptions in the revised version, using the extra page allowed for the camera ready version.

Grokking with attention explained by the proposed theory

We agree that applying our theoretical framework to analyze grokking in Transformers is a promising and worthwhile direction.

As mentioned above, our assumptions are quite generic. They can be shown to hold for attention-based architectures, although establishing Assumption 2 rigorously requires extra work.

However, to show that our results lead to grokking in this case, we need to understand the generalization properties of minimal-norm solutions for Transformers. In Section 5, we focus on simple neural network architectures, for which there already exists a rich body of work characterizing implicit bias and minimal-norm interpolators. Extending this analysis to attention-based models is a highly non-trivial task that calls for a study of its own.

Is the transition at $t=\frac{1}{\lambda}$ theoretically justified? Is the transition sharp or gradual?

Yes, our work gives a precise statement showing that the transitions happens at $t=\frac{1}{\lambda}$. This follows from the change of timescale for $\tilde{w}$, as defined on line 208: we will make this clearer in the revised version.

Indeed, Proposition 2 implies that at time $\varepsilon/\lambda$, the regularized flow remains close to the limit of the unregularized flow, which interpolates but does not generalize well in the high-norm initialization regime. In contrast, at time $\Theta(1/\lambda)$, the flow approaches the limit of the Riemannian gradient flow on the interpolation manifold. This point has smaller norm and, as a result, is expected to generalize better.

As explained in our answer to Reviewer CjUp, this transition is gradual (both theoretically and empirically), but appears sharp when time is plotted in log-scale. Though mathematically incorrect (in the sense that the transition is not “immediate” even in log scale), people tend to say it looks sharp. We claim this is mostly because relatively to the long plateau before the transition, the transition indeed looks “rapid”, i.e. “sharp”.

How small should the weight decay be to observe grokking?

Indeed, our results are derived in the asymptotic regime $\lambda \to 0$, since this allows for a tractable analysis. Extending the theory to a fixed $\lambda > 0$ is considerably harder. That said, we can offer some heuristic analysis regarding how small $\lambda$ needs to be for grokking to emerge.

Grokking depends on a clear separation between two phases: an initial phase where the training loss rapidly converges to zero, and a second, slower phase driven by weight decay, during which test performance improves. For grokking to be observable, the gradient flow should approach a near-stationary point before weight decay begins to significantly influence the dynamics.

To formalize this intuition, we can define two characteristic times: $t_{GF}$​, the convergence time of the unregularized gradient flow, measured as the time at which the gradient norm substantially decreases relative to its initial value; and a second timescale of order $1/\lambda$, associated with the onset of regularization effects. When $\lambda$ is small enough that $t_{GF}\ll 1/\lambda$, we expect to observe grokking-like behavior. Specifically, for some threshold $\varepsilon \ll 1$ (e.g., $\varepsilon=0.01$): $\Vert \nabla F(w_{t_{GF}})\Vert \approx \varepsilon \Vert \nabla F(w_0)\Vert$. Now let $t_{WD}$ denote the time when weight decay kicks in: i.e. when the magnitude of the unregularised gradient becomes comparable to the magnitude of the weight decay term: $\Vert \nabla F(w_{t_{WD}})\Vert \approx \lambda \Vert w_{t_{WD}}\Vert$. Since at time $t_{WD}$, the solution is close to the gradient flow solution $w^{GF}$, we can consider $\Vert w_{t_{WD}}\Vert \approx \Vert w^{GF}\Vert$. The condition for grokking to occur (i.e., a plateau in test loss followed by improvement) is thus that $t_{GF} \ll t_{WD}$. Translating this condition in terms of gradients, we obtain: $\Vert \nabla F(w_{t_{GF}})\Vert \gg \Vert \nabla F(w_{t_{WD}})\Vert$, which, using the approximations above, implies: $\varepsilon \Vert \nabla F(w_0)\Vert \gg \lambda \Vert w^{GF}\Vert$. Simplifying further (absorbing $\varepsilon$ into a constant), we have the practical guideline: $\lambda \ll \frac{\Vert \nabla F(w_0)\Vert }{\Vert w^{GF}\Vert }$. Hence, grokking occurs when weight decay is sufficiently small relative to ratio of the initial gradient magnitude over the norm of the unregularised gradient flow solution.

Empirically, we observed the grokking effect in our experiments for values of $\lambda$ up to $10^{-2}$. For such relatively large values, the initial plateau becomes noticeably shorter. As $\lambda$ increases further, the two phases begin to slightly overlap, making the separation between them less distinct (but still present).

Question on feature unlocking and de-grokking

Is this the “grokking” phenomenon that “unlocks” certain features of the model that would not be possible without “large weight initialization & grokking”, or is it just the problem-intrinsic feature that can also be achieved by de-grokked models?

Our interpretation of the grokking phenomenon is as follows: when initialized with large weights and trained without weight decay, the model does not necessarily learn meaningful features in the sense of prioritizing the relevant ones. Instead, its output depends on a broad set of features, including many that are not informative. In contrast, when regularization is introduced (or in the so-called "rich" regime), we do not believe that the model “unlocks” entirely new features that were absent in the earlier regime. Rather, it learns to concentrate weight primarily on the relevant features–those that contribute meaningfully to generalization–while down-weighting or discarding the irrelevant ones. This contrasts with the unregularized setting (with large initialization), where the model tends to distribute weight more uniformly across both relevant and irrelevant features, with the latter possibly dominating due to their greater frequency.

In this context, the notion of "sparsity" refers to the model ultimately placing significant weight on only a small subset of useful features. Prior to grokking, by contrast, the model assigns non-negligible weights to a broader set of features, many of which are not essential for generalization.